open Test_harness

let position line col offset = { Location.line = line; col = col; offset = offset }

let span line1 col1 line2 col2 =
  Location.make (position line1 col1 0) (position line2 col2 10)

let location_cases =
  let open Location in
  [
    ( "merge of two spans uses first start and second stop",
      fun () ->
        check_equal_string "merged span" "1:2-3:4"
          (to_string (merge (span 1 2 1 5) (span 3 1 3 4))) );
    ("merge with none returns the other span", fun () ->
      check_equal_string "none merge" "2:3-2:3" (to_string (merge none (point (position 2 3 7)))));
    ("is_none distinguishes the empty span", fun () ->
      check "none is none" (is_none none);
      check "point is not none" (not (is_none (point (position 1 1 0)))));
    ("size measures offsets", fun () ->
      check_equal_int "size" 10 (size (span 1 1 2 1)));
    ("pos_to_string renders line and column", fun () ->
      check_equal_string "pos" "12:34" (pos_to_string (position 12 34 100)));
  ]

let ident_cases =
  [
    ( "fresh identifiers get distinct stamps",
      fun () ->
        Ident.reset ();
        let a = Ident.fresh "x" Location.none in
        let b = Ident.fresh "x" Location.none in
        check "distinct" (not (Ident.equal a b));
        check_equal_int "compare is stamp order" (-1) (Ident.compare a b) );
    ( "display keeps the source name while to_string is unique",
      fun () ->
        Ident.reset ();
        let a = Ident.fresh "row" Location.none in
        check_equal_string "display" "row" (Ident.display a);
        check_equal_string "to_string" "row#1" (Ident.to_string a) );
    ( "reset restarts numbering deterministically",
      fun () ->
        Ident.reset ();
        let first = Ident.fresh "x" Location.none in
        Ident.reset ();
        let second = Ident.fresh "x" Location.none in
        check_equal_int "same stamp after reset" (Ident.stamp first) (Ident.stamp second) );
    ( "identifiers carry their source span",
      fun () ->
        Ident.reset ();
        let a = Ident.fresh "x" (span 4 5 4 6) in
        check_equal_string "span" "4:5-4:6" (Location.to_string (Ident.span a)) );
  ]

let diagnostic_cases =
  [
    ( "diagnostic renders file line column and message",
      fun () ->
        let d = Diagnostic.make (span 2 3 2 4) "unknown variable total" in
        check_equal_string "to_string" "<none>:2:3: error: unknown variable total"
          (Diagnostic.to_string d) );
    ( "diagnostic renders a source excerpt with carets",
      fun () ->
        let d = Diagnostic.make (span 2 5 2 9) "bad field" in
        let rendered = Diagnostic.render (Some "let a = 1\nlet b = a.bad\n") d in
        check_true "rendered contains header" rendered
          (rendered = "<none>:2:5: error: bad field\n  let b = a.bad\n      ^^^^") );
    ( "error raises with the current file",
      fun () ->
        Diagnostic.set_file "sample.delta";
        (try
           Diagnostic.error (span 1 1 1 2) "boom";
           fail "raises" "expected an error"
         with Diagnostic.Error d ->
           check_equal_string "file" "sample.delta" d.Diagnostic.file);
        Diagnostic.set_file "" );
  ]

let parse text = Parse.expression text

let parse_error text =
  try
    ignore (parse text);
    None
  with Diagnostic.Error diagnostic -> Some diagnostic

let parse_cases =
  [
    ( "integer and operator spans are recorded",
      fun () ->
        let parsed = parse "1 + 20" in
        check_equal_string "whole span" "1:1-1:7" (Location.to_string parsed.Syntax.espan);
        match parsed.Syntax.e with
        | Syntax.EBinop (Syntax.Add, left, right) ->
            check_equal_string "left span" "1:1-1:2" (Location.to_string left.Syntax.espan);
            check_equal_string "right span" "1:5-1:7" (Location.to_string right.Syntax.espan)
        | _ -> fail "shape" "expected an addition" );
    ( "pipeline is right associative application",
      fun () ->
        match (parse "x |> f |> g").Syntax.e with
        | Syntax.EApp (outer, inner) -> (
            check "outer argument is the inner application"
              (match outer.Syntax.e with Syntax.EVar "g" -> true | _ -> false);
            match inner.Syntax.e with
            | Syntax.EApp (f, arg) ->
                check "inner function is f" (match f.Syntax.e with Syntax.EVar "f" -> true | _ -> false);
                check "source is x" (match arg.Syntax.e with Syntax.EVar "x" -> true | _ -> false)
            | _ -> fail "shape" "expected an application on the left" )
        | _ -> fail "shape" "expected an application" );
    ( "let, lambda and conditional parse with spans",
      fun () ->
        match (parse "let f = fun n -> if n > 0 then n else - n in f 3").Syntax.e with
        | Syntax.ELet (name, _, body) ->
            check_equal_string "bound name" "f" name;
            check "body is an application" (match body.Syntax.e with Syntax.EApp _ -> true | _ -> false)
        | _ -> fail "shape" "expected a let expression" );
    ( "tuples and record insertion parse",
      fun () ->
        match (parse "(1, 2, 3)").Syntax.e with
        | Syntax.ETuple [ _; _; _ ] -> check "tuple of three" true
        | _ -> fail "shape" "expected a tuple";
        match (parse "{ customer = \"Ada\" }").Syntax.e with
        | Syntax.ERecord [ ("customer", _) ] -> check "record with one field" true
        | _ -> fail "shape" "expected a record" );
    ( "field projection binds tighter than application",
      fun () ->
        match (parse "f o.total").Syntax.e with
        | Syntax.EApp (_, argument) ->
            check "argument is a projection" (match argument.Syntax.e with Syntax.EField _ -> true | _ -> false)
        | _ -> fail "shape" "expected an application" );
    ( "comparison binds looser than conjunction",
      fun () ->
        match (parse "a && b = c").Syntax.e with
        | Syntax.EBinop (Syntax.And, _, right) ->
            check "right side is a comparison" (match right.Syntax.e with Syntax.EBinop (Syntax.Eq, _, _) -> true | _ -> false)
        | _ -> fail "shape" "expected a conjunction at the root" );
    ( "unexpected character reports a span",
      fun () ->
        match parse_error "let x = @" with
        | None -> fail "lexical error" "expected a diagnostic"
        | Some diagnostic ->
            check_equal_string "message" "<none>:1:9: error: unexpected character '@'"
              (Diagnostic.to_string diagnostic) );
    ( "unterminated string reports the opening quote",
      fun () ->
        match parse_error "let s = \"abc" with
        | None -> fail "string" "expected a diagnostic"
        | Some diagnostic ->
            check_equal_string "message" "<none>:1:9: error: unterminated string literal"
              (Diagnostic.to_string diagnostic);
            check_equal_string "span starts at the opening quote" "1:9-1:13"
              (Location.to_string diagnostic.Diagnostic.span) );
    ( "unterminated comment reports a span",
      fun () ->
        match parse_error "1 (* comment" with
        | None -> fail "comment" "expected a diagnostic"
        | Some diagnostic ->
            check_equal_string "message" "<none>:1:3: error: unterminated comment"
              (Diagnostic.to_string diagnostic);
            check_equal_string "span starts at the comment opener" "1:3-1:13"
              (Location.to_string diagnostic.Diagnostic.span) );
    ( "integer literal overflow reports a span",
      fun () ->
        match parse_error "99999999999999999999999" with
        | None -> fail "overflow" "expected a diagnostic"
        | Some diagnostic ->
            check "message is non empty" (String.length diagnostic.Diagnostic.message > 0);
            check_equal_string "position" "1:1" (Location.pos_to_string diagnostic.Diagnostic.span.Location.start) );
    ( "missing else reports the offending token",
      fun () ->
        match parse_error "if a then b" with
        | None -> fail "syntax" "expected a diagnostic"
        | Some diagnostic ->
            check_equal_string "message" "<none>:1:12: error: syntax error: unexpected end of file"
              (Diagnostic.to_string diagnostic) );
    ( "nested comments are skipped",
      fun () ->
        let parsed = parse "(* a (* b *) c *) 7" in
        check "the expression after the comment is parsed"
          (match parsed.Syntax.e with Syntax.EInt 7 -> true | _ -> false) );
  ]

let program_error text =
  try
    ignore (Parse.program text);
    None
  with Diagnostic.Error diagnostic -> Some diagnostic

let program_cases =
  [
    ( "the required example parses into declarations",
      fun () ->
        let parsed = Parse.program (read_fixture "expensive_order.delta") in
        check_equal_int "one record declaration" 1 (List.length parsed.Syntax.prog_records);
        let record = List.hd parsed.Syntax.prog_records in
        check_equal_string "record name" "order" record.Syntax.rd_name;
        check_equal_int "two fields" 2 (List.length record.Syntax.rd_fields);
        check_equal_string "input name" "orders" parsed.Syntax.prog_input.Syntax.in_name;
        check_equal_string "input type" "collection order"
          (Syntax.string_of_tyexpr parsed.Syntax.prog_input.Syntax.in_ty);
        check_equal_int "one helper" 1 (List.length parsed.Syntax.prog_helpers);
        let helper = List.hd parsed.Syntax.prog_helpers in
        check_equal_string "helper name" "tax" helper.Syntax.h_name;
        check_equal_int "one parameter" 1 (List.length helper.Syntax.h_params);
        check_equal_string "query name" "expensive_orders" parsed.Syntax.prog_query.Syntax.q_name );
    ( "a query pipeline of filter and map is right nested",
      fun () ->
        let parsed = Parse.program (read_fixture "expensive_order.delta") in
        match parsed.Syntax.prog_query.Syntax.q_body.Syntax.e with
        | Syntax.EApp (mapper, upstream) ->
            check "outer operator is map" (match mapper.Syntax.e with Syntax.EMap _ -> true | _ -> false);
            check "upstream is a filter" (match upstream.Syntax.e with Syntax.EApp (f, _) -> (match f.Syntax.e with Syntax.EFilter _ -> true | _ -> false) | _ -> false)
        | _ -> fail "shape" "expected a pipeline application" );
    ( "sum and count parse as collection operators",
      fun () ->
        let parsed = Parse.program (read_fixture "revenue.delta") in
        check "sum query is an application of sum"
          (match parsed.Syntax.prog_query.Syntax.q_body.Syntax.e with
          | Syntax.EApp (sum, _) -> (match sum.Syntax.e with Syntax.ESum -> true | _ -> false)
          | _ -> false);
        let parsed = Parse.program (read_fixture "count_large.delta") in
        check "count query is an application of count"
          (match parsed.Syntax.prog_query.Syntax.q_body.Syntax.e with
          | Syntax.EApp (count, _) -> (match count.Syntax.e with Syntax.ECount -> true | _ -> false)
          | _ -> false) );
    ( "a trailing semicolon in a record declaration is accepted",
      fun () ->
        let parsed =
          Parse.program "type row = { a : int; b : string; }\ninput rows : collection row\nquery q = rows\n"
        in
        check_equal_int "two fields" 2 (List.length (List.hd parsed.Syntax.prog_records).Syntax.rd_fields) );
    ( "declarations may appear in any order",
      fun () ->
        let parsed =
          Parse.program "query q = rows\nlet f n = n\ninput rows : collection row\ntype row = { a : int }\n"
        in
        check_equal_string "input" "rows" parsed.Syntax.prog_input.Syntax.in_name );
    ( "a program without an input is rejected",
      fun () ->
        match program_error "query q = 1\n" with
        | None -> fail "input" "expected a diagnostic"
        | Some diagnostic ->
            check_equal_string "message"
              "the program declares no input collection: add `input NAME : collection T`"
              diagnostic.Diagnostic.message );
    ( "a program without a query is rejected",
      fun () ->
        match program_error "input rows : collection row\ntype row = { a : int }\n" with
        | None -> fail "query" "expected a diagnostic"
        | Some diagnostic ->
            check_equal_string "message" "<none>: error: the program declares no query: add `query NAME = EXPRESSION`"
              (Diagnostic.to_string diagnostic) );
    ( "a second input declaration points at the offending one",
      fun () ->
        match program_error "input a : collection int\ninput b : collection int\nquery q = a\n" with
        | None -> fail "input" "expected a diagnostic"
        | Some diagnostic ->
            check_equal_string "message" "program.delta:2:1: error: the program declares more than one input collection"
              (Diagnostic.to_string { diagnostic with Diagnostic.file = "program.delta" }) );
    ( "a malformed record declaration reports a span",
      fun () ->
        match program_error "type row = { a int }\n" with
        | None -> fail "record" "expected a diagnostic"
        | Some diagnostic ->
            check "diagnostic points at line 1"
              (diagnostic.Diagnostic.span.Location.start.Location.line = 1);
            check "message is non empty" (String.length diagnostic.Diagnostic.message > 0) );
  ]

let resolve_program text = Resolve.program (Parse.program text)

let resolve_error text =
  try
    ignore (resolve_program text);
    None
  with Diagnostic.Error diagnostic -> Some diagnostic

let resolve_cases =
  [
    ( "the required example resolves to fresh identifiers",
      fun () ->
        Ident.reset ();
        let resolved = resolve_program (read_fixture "expensive_order.delta") in
        check_equal_string "input binding" "orders" (Ident.name resolved.Resolve.rp_input);
        check_equal_int "input stamp" 1 (Ident.stamp resolved.Resolve.rp_input) );
    ( "shadowing binds a new identifier",
      fun () ->
        Ident.reset ();
        let resolved = resolve_program "input rows : collection int\nquery q = let x = 1 in let x = x + 1 in x\n" in
        let rec collect expr acc =
          match expr.Resolve.r with
          | Resolve.RLet (ident, bound, body) -> collect bound (ident :: acc) |> fun acc -> collect body acc
          | _ -> acc
        in
        let binders = collect resolved.Resolve.rp_query_body [] in
        match binders with
        | [ outer; inner ] ->
            check_equal_string "both binders are called x" "x" (Ident.name inner);
            check "shadowed binders differ" (Ident.stamp outer <> Ident.stamp inner)
        | _ -> fail "shadowing" "expected two let binders" );
    ( "an unbound variable is reported with its span",
      fun () ->
        match
          resolve_error "input rows : collection int\nquery q = rows |> filter (fun r -> totl > 1)\n"
        with
        | None -> fail "unbound" "expected a diagnostic"
        | Some diagnostic ->
            check_equal_string "message" "unbound variable `totl`" diagnostic.Diagnostic.message;
            check_equal_string "span" "2:36-2:40" (Location.to_string diagnostic.Diagnostic.span) );
    ( "the input collection is rejected inside a helper",
      fun () ->
        match resolve_error "input rows : collection int\nlet total = sum rows\nquery q = rows\n" with
        | None -> fail "input use" "expected a diagnostic"
        | Some diagnostic ->
            check_equal_string "message"
              "the input collection `rows` may only be used in the query, not inside a helper function"
              diagnostic.Diagnostic.message );
    ( "duplicate helper names are rejected",
      fun () ->
        match resolve_error "input rows : collection int\nlet f n = n\nlet f n = n + 1\nquery q = rows\n" with
        | None -> fail "duplicate" "expected a diagnostic"
        | Some diagnostic ->
            check_equal_string "message" "duplicate helper function `f`, first declared at 2:1-2:12"
              diagnostic.Diagnostic.message );
    ( "duplicate record names are rejected",
      fun () ->
        match resolve_error "type r = { a : int }\ntype r = { b : int }\ninput rows : collection r\nquery q = rows\n" with
        | None -> fail "duplicate" "expected a diagnostic"
        | Some diagnostic ->
            check "message is non empty" (String.length diagnostic.Diagnostic.message > 0);
            check_equal_string "line" "2"
              (string_of_int diagnostic.Diagnostic.span.Location.start.Location.line) );
    ( "field labels must be globally unique",
      fun () ->
        match resolve_error "type a = { x : int }\ntype b = { x : int }\ninput rows : collection a\nquery q = rows\n" with
        | None -> fail "labels" "expected a diagnostic"
        | Some diagnostic ->
            check_equal_string "message"
              "duplicate field label `x`, already declared in record `a`; field labels must be globally unique"
              diagnostic.Diagnostic.message );
    ( "recursive helpers are rejected",
      fun () ->
        match resolve_error "input rows : collection int\nlet f n = f n\nquery q = rows\n" with
        | None -> fail "recursion" "expected a diagnostic"
        | Some diagnostic ->
            check_equal_string "message"
              "helper function `f` is recursive; recursive definitions are not supported"
              diagnostic.Diagnostic.message );
    ( "mutually recursive helpers are rejected",
      fun () ->
        match resolve_error "input rows : collection int\nlet f n = g n\nlet g n = f n\nquery q = rows\n" with
        | None -> fail "recursion" "expected a diagnostic"
        | Some diagnostic ->
            check "message is non empty" (String.length diagnostic.Diagnostic.message > 0) );
    ( "duplicate parameters are rejected",
      fun () ->
        match resolve_error "input rows : collection int\nlet f n n = n\nquery q = rows\n" with
        | None -> fail "parameters" "expected a diagnostic"
        | Some diagnostic -> check_equal_string "message" "duplicate parameter `n`" diagnostic.Diagnostic.message );
    ( "helpers may call other helpers",
      fun () ->
        let resolved =
          resolve_program
            "input rows : collection int\nlet double n = n + n\nlet quad n = double (double n)\nquery q = rows |> map (fun r -> quad r)\n"
        in
        check_equal_int "two helpers" 2 (List.length resolved.Resolve.rp_helpers) );
  ]

let () =
  Test_harness.run_suite "location" location_cases;
  Test_harness.run_suite "ident" ident_cases;
  Test_harness.run_suite "diagnostic" diagnostic_cases;
  Test_harness.run_suite "parse" parse_cases;
  Test_harness.run_suite "program" program_cases;
  Test_harness.run_suite "resolve" resolve_cases;
  Test_harness.run_suite "types" Test_type.cases;
  Test_harness.run_suite "specialize" Test_type.specialize_cases;
  Test_harness.run_suite "anf" Test_type.anf_cases;
  Test_harness.run_suite "changes" Test_change.law_cases;
  Test_harness.run_suite "batches" Test_change.batch_cases;
  Test_harness.run_suite "interpret" Test_incremental.cases;
  Test_harness.run_suite "graph" Test_incremental.graph_cases;
  Test_harness.run_suite "executor" Test_incremental.executor_cases;
  Test_harness.run_suite "filter" Test_incremental.filter_cases;
  Test_harness.run_suite "aggregates" Test_incremental.aggregate_cases;
  Test_harness.run_suite "simplify" Test_incremental.simplify_cases;
  Test_harness.run_suite "codegen" Test_codegen.codegen_cases;
  Test_harness.run_suite "generated updates" Test_codegen.update_cases;
  Test_harness.run_suite "wire" Test_codegen.wire_cases;
  Test_harness.run_suite "cli" Test_codegen.cli_cases;
  Test_harness.run_suite "install" Test_codegen.install_cases;
  Test_harness.run_suite "regression" Test_incremental.regression_cases;
  Test_harness.run_suite "differential" Test_incremental.differential_cases;
  Test_harness.run_suite "generated differential" Test_codegen.generated_differential_cases;
  Printf.printf "%d cases, %d failures\n" (Test_harness.case_count ()) (Test_harness.failure_count ());
  exit (if Test_harness.failure_count () = 0 then 0 else 1)