open Test_harness

let row customer total = Value.VRecord ("order", [ ("customer", Value.VString customer); ("total", Value.VInt total) ])

let input entries = List.mapi (fun index entry -> (index + 1, entry)) entries

let run fixture_name entries =
  let program = infer (read_fixture fixture_name) in
  Interpret.program program (input entries)

let run_text text entries =
  let program = infer text in
  Interpret.program program (input entries)

let cases =
  [
    ( "the example query selects and projects rows",
      fun () ->
        let result =
          run "expensive_order.delta"
            [ row "Ada" 1500; row "Bo" 900; row "Lin" 2200 ]
        in
        check_equal_string "result" "(collection (1 (tuple \"Ada\" 300)) (3 (tuple \"Lin\" 440)))"
          (Value.to_string result) );
    ( "a filter that keeps nothing yields an empty collection",
      fun () ->
        let result = run "expensive_order.delta" [ row "Bo" 10 ] in
        check_equal_string "empty" "(collection )" (Value.to_string result) );
    ( "an empty input yields an empty collection and zero aggregates",
      fun () ->
        check_equal_string "no rows" "(collection )" (Value.to_string (run "expensive_order.delta" []));
        check_equal_int "revenue" 0
          (match run "revenue.delta" [] with Value.VInt total -> total | _ -> -1);
        check_equal_int "count" 0
          (match run "count_large.delta" [] with Value.VInt total -> total | _ -> -1) );
    ( "revenue sums the taxed totals",
      fun () ->
        let result = run "revenue.delta" [ row "Ada" 1000; row "Bo" 250 ] in
        check_equal_string "revenue" "250" (Value.to_string result) );
    ( "count_large counts the retained rows",
      fun () ->
        let result = run "count_large.delta" [ row "Ada" 1000; row "Bo" 250; row "Lin" 501 ] in
        check_equal_string "count" "2" (Value.to_string result) );
    ( "negative values are preserved",
      fun () ->
        let result =
          run_text "input rows : collection int\nquery q = rows |> filter (fun r -> r < 0) |> sum\n"
            [ Value.VInt (-5); Value.VInt 3; Value.VInt (-7) ]
        in
        check_equal_string "sum" "-12" (Value.to_string result) );
    ( "integer division truncates towards zero",
      fun () ->
        let result = run_text "input rows : collection int\nquery q = rows |> map (fun r -> r / 3) |> sum\n"
            [ Value.VInt 10; Value.VInt (-10) ]
        in
        check_equal_string "sum" "0" (Value.to_string result) );
    ( "division by zero is a runtime error",
      fun () ->
        (try
           ignore (run_text "input rows : collection int\nquery q = rows |> map (fun r -> 100 / r) |> sum\n"
                     [ Value.VInt 5; Value.VInt 0 ]);
           fail "division" "expected a runtime error"
         with Division_by_zero -> check "raised" true) );
    ( "the mapping expression runs only where the filter keeps rows",
      fun () ->
        let result =
          run_text "input rows : collection int\nquery q = rows |> filter (fun r -> r > 10) |> map (fun r -> 1000 / r) |> sum\n"
            [ Value.VInt 5; Value.VInt 20 ]
        in
        check_equal_string "sum" "50" (Value.to_string result) );
    ( "boolean operators short circuit",
      fun () ->
        let result =
          run_text "input rows : collection int\nquery q = rows |> count\n"
            [ Value.VInt 1 ]
        in
        check_equal_string "count" "1" (Value.to_string result);
        let result =
          run_text "input rows : collection int\nquery q = rows |> filter (fun r -> false && 1 / r > 0) |> count\n"
            [ Value.VInt 0 ]
        in
        check_equal_string "short circuit and" "0" (Value.to_string result);
        let result =
          run_text "input rows : collection int\nquery q = rows |> filter (fun r -> true || 1 / r > 0) |> count\n"
            [ Value.VInt 0 ]
        in
        check_equal_string "short circuit or" "1" (Value.to_string result) );
    ( "map preserves keys and filter keeps the original keys",
      fun () ->
        let result =
          run_text "input rows : collection int\nquery q = rows |> filter (fun r -> r > 1) |> map (fun r -> r * 10)\n"
            [ Value.VInt 1; Value.VInt 2; Value.VInt 3 ]
        in
        check_equal_string "keys" "(collection (2 20) (3 30))" (Value.to_string result) );
    ( "helpers are applied at their call sites",
      fun () ->
        let result =
          run_text
            "input rows : collection int\nlet scale n = n * 3\nlet offset n = scale n + 1\nquery q = rows |> map offset |> sum\n"
            [ Value.VInt 1; Value.VInt 2 ]
        in
        check_equal_string "sum" "11" (Value.to_string result) );
    ( "records and tuples are compared structurally",
      fun () ->
        let result =
          run_text
            "type pair = { first : int; second : int }\ninput rows : collection pair\nquery q = rows |> filter (fun r -> r = { first = 1; second = 2 }) |> count\n"
            [ Value.VRecord ("pair", [ ("first", Value.VInt 1); ("second", Value.VInt 2) ]); Value.VRecord ("pair", [ ("first", Value.VInt 2); ("second", Value.VInt 1) ]) ]
        in
        check_equal_string "count" "1" (Value.to_string result) );
    ( "a constant query ignores the input",
      fun () ->
        let result = run_text "input rows : collection int\nquery q = 6 * 7\n" [ Value.VInt 1 ] in
        check_equal_string "constant" "42" (Value.to_string result) );
    ( "out of range arithmetic follows machine integers",
      fun () ->
        let result =
          run_text "input rows : collection int\nquery q = rows |> sum\n"
            [ Value.VInt max_int; Value.VInt 1 ]
        in
        check_equal_string "wrapped" (string_of_int min_int) (Value.to_string result) );
  ]

let plan text = Graph.build (Anf.program (Specialize.program (infer text)))

let plan_of_fixture name = Graph.build (Anf.program (Specialize.program (infer (read_fixture name))))

let graph_cases =
  [
    ( "the example plan is a source, a filter and a map",
      fun () ->
        let plan = plan_of_fixture "expensive_order.delta" in
        check_equal_int "three nodes" 3 (List.length plan.Graph.pl_nodes);
        let nodes = plan.Graph.pl_nodes in
        check "source first" (match (List.nth nodes 0).Graph.n_kind with Graph.Source -> true | _ -> false);
        check "filter second"
          (match (List.nth nodes 1).Graph.n_kind with Graph.Filter _ -> true | _ -> false);
        check "map last" (match (List.nth nodes 2).Graph.n_kind with Graph.Map _ -> true | _ -> false);
        check_equal_string "output" "collection (string, int)" (Types.pp plan.Graph.pl_output) );
    ( "the revenue plan ends in a sum accumulator",
      fun () ->
        let plan = plan_of_fixture "revenue.delta" in
        let root = Graph.node_of_id plan (match plan.Graph.pl_result with Graph.Result_collection id -> id | _ -> -1) in
        check "sum" (match root.Graph.n_kind with Graph.Sum -> true | _ -> false);
        check "accumulator cache" (root.Graph.n_cache = Graph.Accumulator);
        check_equal_string "linear output" "int" (Types.pp plan.Graph.pl_output) );
    ( "the count query counts the retained rows",
      fun () ->
        let plan = plan_of_fixture "count_large.delta" in
        check_equal_int "three nodes" 3 (List.length plan.Graph.pl_nodes);
        let root =
          match plan.Graph.pl_result with
          | Graph.Result_collection id -> Graph.node_of_id plan id
          | Graph.Result_scalar _ ->
              fail "plan" "expected a node result";
              { Graph.n_id = -1; n_kind = Graph.Source; n_input = None; n_element = Types.TInt; n_cache = Graph.No_cache; n_span = Location.none }
        in
        check "count" (match root.Graph.n_kind with Graph.Count -> true | _ -> false) );
    ( "collection nodes record their consumers",
      fun () ->
        let plan = plan_of_fixture "expensive_order.delta" in
        check_equal_string "source consumers" "1" (Util.join "," (List.map string_of_int plan.Graph.pl_consumers.(0)));
        check_equal_string "filter consumers" "2" (Util.join "," (List.map string_of_int plan.Graph.pl_consumers.(1)));
        check_equal_string "map has no consumers" "" (Util.join "," (List.map string_of_int plan.Graph.pl_consumers.(2))) );
    ( "an identity query has a single source node",
      fun () ->
        let plan = plan "input rows : collection int\nquery q = rows\n" in
        check_equal_int "one node" 1 (List.length plan.Graph.pl_nodes);
        check "the result is the source"
          (match plan.Graph.pl_result with Graph.Result_collection 0 -> true | _ -> false) );
    ( "a plan dump is deterministic",
      fun () ->
        Ident.reset ();
        Types.reset ();
        Anf.reset ();
        let first = Graph.dump (plan_of_fixture "expensive_order.delta") in
        Ident.reset ();
        Types.reset ();
        Anf.reset ();
        let second = Graph.dump (plan_of_fixture "expensive_order.delta") in
        check_equal_string "identical" first second );
    ( "an integer query that only uses the input is a plain plan",
      fun () ->
        let plan = plan "input rows : collection int\nquery q = rows |> sum\n" in
        check_equal_int "two nodes" 2 (List.length plan.Graph.pl_nodes);
        check "integer output" (Types.repr plan.Graph.pl_output = Types.TInt) );
    ( "a constant integer query produces no collection nodes",
      fun () ->
        let plan = plan "input rows : collection int\nquery q = 40 + 2\n" in
        check_equal_int "no nodes" 0 (List.length plan.Graph.pl_nodes);
        check "scalar result" (match plan.Graph.pl_result with Graph.Result_scalar _ -> true | _ -> false) );
    ( "shared collection temporaries are reuse",
      fun () ->
        let plan =
          plan
            "input rows : collection int\nquery q = rows |> filter (fun r -> r > 0) |> map (fun r -> r + 1) |> sum\n"
        in
        check_equal_int "four nodes" 4 (List.length plan.Graph.pl_nodes) );
  ]

type fixture = {
  fx_program : Typed.program;
  fx_plan : Graph.plan;
  fx_counters : Delta_runtime.counters;
  mutable fx_state : Incremental.state;
  fx_entries : (int * Value.t) list;
}

let build_fixture text entries =
  let program = infer text in
  let plan = Graph.build (Anf.program (Specialize.program program)) in
  let counters = Delta_runtime.new_counters () in
  let state = Incremental.init plan counters entries in
  { fx_program = program; fx_plan = plan; fx_counters = counters; fx_state = state; fx_entries = entries }

let fixture_entries plan state = Delta_runtime.Pure_map.bindings state.Incremental.s_input

let reference_result fixture state =
  Interpret.program fixture.fx_program (fixture_entries fixture.fx_plan state)

let reference_after fixture ops =
  match Change.validate_batch ~existing:fixture.fx_state.Incremental.s_input ops with
  | Change.Failure message -> Error message
  | Change.Success (temp, _) ->
      let program = fixture.fx_program in
      ignore program;
      Ok (Interpret.program fixture.fx_program (Delta_runtime.Pure_map.bindings temp))

let step fixture ops =
  let before = Incremental.result fixture.fx_plan fixture.fx_state in
  match Incremental.apply_batch fixture.fx_plan fixture.fx_counters fixture.fx_state ops with
  | Change.Failure message -> Error message
  | Change.Success (state, change) ->
      let applied = Change.apply before change in
      let cached_after = Incremental.result fixture.fx_plan state in
      fixture.fx_state <- state;
      Ok (applied, cached_after, change)

let order_row customer total =
  Value.VRecord ("order", [ ("customer", Value.VString customer); ("total", Value.VInt total) ])

let executor_cases =
  let orders text entries =
    let fixture =
      build_fixture
        ("input orders : collection order\n" ^ text)
        entries
    in
    fixture
  in
  [
    ( "initialization matches the reference interpreter",
      fun () ->
        let fixture =
          build_fixture (read_fixture "expensive_order.delta")
            [ (1, order_row "Ada" 1500); (2, order_row "Bo" 900) ]
        in
        check_equal_string "initial result"
          (Value.to_string (reference_result fixture fixture.fx_state))
          (Value.to_string (Incremental.result fixture.fx_plan fixture.fx_state)) );
    ( "inserting a key appends a mapped contribution",
      fun () ->
        let fixture =
          build_fixture (read_fixture "expensive_order.delta") [ (1, order_row "Ada" 1500) ]
        in
        (match step fixture [ Change.OpInsert (2, order_row "Lin" 2200) ] with
        | Error message -> fail "insert" message
        | Ok (applied, cached, change) ->
            let expected = reference_result fixture fixture.fx_state in
            check_equal_string "change" "collection [insert 2 (tuple \"Lin\" 440)]"
              (Change.to_string change);
            check_equal_string "applied" (Value.to_string expected) (Value.to_string applied);
            check_equal_string "cached" (Value.to_string expected) (Value.to_string cached)) );
    ( "removing a key retracts the cached contribution",
      fun () ->
        let fixture =
          build_fixture (read_fixture "expensive_order.delta")
            [ (1, order_row "Ada" 1500); (2, order_row "Lin" 2200) ]
        in
        (match step fixture [ Change.OpRemove 2 ] with
        | Error message -> fail "remove" message
        | Ok (applied, cached, change) ->
            let expected = reference_result fixture fixture.fx_state in
            check_equal_string "change" "collection [remove 2 (tuple \"Lin\" 440)]"
              (Change.to_string change);
            check_equal_string "applied" (Value.to_string expected) (Value.to_string applied);
            check_equal_string "cached" (Value.to_string expected) (Value.to_string cached)) );
    ( "replacing a retained key forwards the new contribution",
      fun () ->
        let fixture =
          build_fixture (read_fixture "expensive_order.delta") [ (1, order_row "Ada" 1500) ]
        in
        (match step fixture [ Change.OpReplace (1, order_row "Ada" 2500) ] with
        | Error message -> fail "replace" message
        | Ok (applied, cached, change) ->
            let expected = reference_result fixture fixture.fx_state in
            check_equal_string "change" "collection [replace 1 (tuple \"Ada\" 300) (tuple \"Ada\" 500)]"
              (Change.to_string change);
            check_equal_string "applied" (Value.to_string expected) (Value.to_string applied);
            check_equal_string "cached" (Value.to_string expected) (Value.to_string cached)) );
    ( "replacing a key with an equal row produces no change",
      fun () ->
        let fixture =
          build_fixture (read_fixture "expensive_order.delta") [ (1, order_row "Ada" 1500) ]
        in
        (match step fixture [ Change.OpReplace (1, order_row "Ada" 1500) ] with
        | Error message -> fail "equal" message
        | Ok (applied, cached, change) ->
            check "empty change" (Change.is_empty change);
            check_equal_string "applied" (Value.to_string (reference_result fixture fixture.fx_state))
              (Value.to_string applied);
            check_equal_string "cached" (Value.to_string (reference_result fixture fixture.fx_state))
              (Value.to_string cached)) );
    ( "a row that stops matching leaves the collection only once",
      fun () ->
        let fixture =
          build_fixture (read_fixture "expensive_order.delta") [ (1, order_row "Ada" 1500) ]
        in
        (match step fixture [ Change.OpReplace (1, order_row "Ada" 100) ] with
        | Error message -> fail "drop" message
        | Ok (applied, cached, change) ->
            let expected = reference_result fixture fixture.fx_state in
            check_equal_string "change" "collection [remove 1 (tuple \"Ada\" 300)]"
              (Change.to_string change);
            check_equal_string "applied" (Value.to_string expected) (Value.to_string applied);
            check_equal_string "cached" (Value.to_string expected) (Value.to_string cached)) );
    ( "a row that starts matching is inserted once",
      fun () ->
        let fixture =
          build_fixture (read_fixture "expensive_order.delta") [ (1, order_row "Ada" 100) ]
        in
        (match step fixture [ Change.OpReplace (1, order_row "Ada" 5000) ] with
        | Error message -> fail "gain" message
        | Ok (applied, cached, change) ->
            let expected = reference_result fixture fixture.fx_state in
            check_equal_string "change" "collection [insert 1 (tuple \"Ada\" 1000)]"
              (Change.to_string change);
            check_equal_string "applied" (Value.to_string expected) (Value.to_string applied);
            check_equal_string "cached" (Value.to_string expected) (Value.to_string cached)) );
    ( "several batches keep the state consistent",
      fun () ->
        let fixture =
          build_fixture (read_fixture "expensive_order.delta") [ (1, order_row "Ada" 1500) ]
        in
        (match step fixture [ Change.OpInsert (2, order_row "Bo" 3000) ] with
        | Error message -> fail "insert" message
        | Ok _ -> ());
        (match step fixture [ Change.OpRemove 1 ] with
        | Error message -> fail "remove" message
        | Ok _ -> ());
        (match step fixture [ Change.OpReplace (2, order_row "Bo" 900) ] with
        | Error message -> fail "replace" message
        | Ok (applied, cached, _) ->
            let expected = reference_result fixture fixture.fx_state in
            check_equal_string "applied" (Value.to_string expected) (Value.to_string applied);
            check_equal_string "cached" (Value.to_string expected) (Value.to_string cached)) );
    ( "an identity query forwards the input change",
      fun () ->
        let fixture =
          build_fixture "input rows : collection int\nquery q = rows" [ (1, Value.VInt 5) ]
        in
        (match step fixture [ Change.OpInsert (2, Value.VInt 9) ] with
        | Error message -> fail "identity" message
        | Ok (applied, cached, change) ->
            check_equal_string "change" "collection [insert 2 9]" (Change.to_string change);
            check_equal_string "applied" "(collection (1 5) (2 9))" (Value.to_string applied);
            check_equal_string "cached" "(collection (1 5) (2 9))" (Value.to_string cached)) );
    ( "an integer query reports an additive change",
      fun () ->
        let fixture =
          build_fixture (read_fixture "revenue.delta") [ (1, order_row "Ada" 1000) ]
        in
        (match step fixture [ Change.OpInsert (2, order_row "Bo" 250) ] with
        | Error message -> fail "sum" message
        | Ok (applied, cached, change) ->
            check_equal_string "change" "+50" (Change.to_string change);
            check_equal_string "applied" "250" (Value.to_string applied);
            check_equal_string "cached" "250" (Value.to_string cached)) );
    ( "count ignores replacements that keep a row",
      fun () ->
        let fixture =
          build_fixture (read_fixture "count_large.delta") [ (1, order_row "Ada" 1000) ]
        in
        (match step fixture [ Change.OpReplace (1, order_row "Ada" 2000) ] with
        | Error message -> fail "count" message
        | Ok (_, _, change) -> check "empty change" (Change.is_empty change));
        (match step fixture [ Change.OpReplace (1, order_row "Ada" 10) ] with
        | Error message -> fail "count" message
        | Ok (_, cached, change) ->
            check_equal_string "change" "-1" (Change.to_string change);
            check_equal_string "cached" "0" (Value.to_string cached)) );
    ( "an invalid batch leaves the state usable",
      fun () ->
        let fixture =
          build_fixture (read_fixture "expensive_order.delta") [ (1, order_row "Ada" 1500) ]
        in
        let before = Value.to_string (Incremental.result fixture.fx_plan fixture.fx_state) in
        (match step fixture [ Change.OpRemove 99 ] with
        | Ok _ -> fail "invalid" "expected a failure"
        | Error _ -> ());
        check_equal_string "result unchanged" before
          (Value.to_string (Incremental.result fixture.fx_plan fixture.fx_state));
        (match step fixture [ Change.OpInsert (2, order_row "Bo" 2000) ] with
        | Error message -> fail "recovery" message
        | Ok (applied, _, _) ->
            check_equal_string "still incremental" (Value.to_string (reference_result fixture fixture.fx_state))
              (Value.to_string applied)) );
    ( "updates visit one key per node",
      fun () ->
        let fixture =
          build_fixture (read_fixture "expensive_order.delta")
            [ (1, order_row "Ada" 1500); (2, order_row "Bo" 2000); (3, order_row "Cy" 500) ]
        in
        Delta_runtime.reset_counters fixture.fx_counters;
        (match step fixture [ Change.OpReplace (2, order_row "Bo" 2500) ] with
        | Error message -> fail "visit" message
        | Ok _ ->
            let counters = fixture.fx_counters in
            check_equal_int "changed key visits" 3 counters.Delta_runtime.changed_key_visits;
            check_equal_int "predicate evaluations" 1 counters.Delta_runtime.predicate_evaluations;
            check_equal_int "mapping evaluations" 0 counters.Delta_runtime.mapping_evaluations;
            check_equal_int "scalar deltas" 1 counters.Delta_runtime.scalar_deltas;
            check_equal_int "full traversals" 0 counters.Delta_runtime.full_traversals) );
    ( "initialization counts one traversal per node",
      fun () ->
        let plan = plan_of_fixture "expensive_order.delta" in
        let counters = Delta_runtime.new_counters () in
        let _ = Incremental.init plan counters [ (1, order_row "Ada" 1500) ] in
        check_equal_int "full traversals" 3 counters.Delta_runtime.full_traversals;
        check_equal_int "mapping evaluations" 1 counters.Delta_runtime.mapping_evaluations;
        check_equal_int "predicate evaluations" 1 counters.Delta_runtime.predicate_evaluations );
  ]