vanity/src/gen.ml

open Types

type specimen = {
  ty : typ;
  term : Source.term;
}

type result = {
  total : int;
  related : int;
  failures : (specimen * Relation.comparison) list;
}

let rng = Random.State.make [| 0x51; 0x17; 0x2b |]

let fresh =
  let c = ref 0 in
  fun prefix ->
    incr c;
    prefix ^ string_of_int !c

let pick xs = List.nth xs (Random.State.int rng (List.length xs))

let rec gen_typ depth =
  if depth = 0 then pick [TInt; TBool]
  else
    pick
      [
        TInt;
        TBool;
        TPair (gen_typ (depth - 1), gen_typ (depth - 1));
        TSum (gen_typ (depth - 1), gen_typ (depth - 1));
      ]

let rec gen_closed depth ty =
  if depth = 0 then
    match ty with
    | TInt -> Source.Int (Random.State.int rng 5)
    | TBool -> Source.Bool (Random.State.bool rng)
    | TPair (a, b) -> Source.Pair (gen_closed 0 a, gen_closed 0 b)
    | TSum (a, b) ->
        if Random.State.bool rng then Source.Inl (TSum (a, b), gen_closed 0 a)
        else Source.Inr (TSum (a, b), gen_closed 0 b)
    | _ -> Source.Int 0
  else
    match ty with
    | TInt ->
        pick
          [
            Source.Int (Random.State.int rng 7);
            Source.If (Source.Bool (Random.State.bool rng), Source.Int 1, Source.Int 0);
          ]
    | TBool ->
        pick
          [
            Source.Bool (Random.State.bool rng);
            Source.Eq (GBool, Source.Bool (Random.State.bool rng), Source.Bool (Random.State.bool rng));
          ]
    | TPair (a, b) -> Source.Pair (gen_closed (depth - 1) a, gen_closed (depth - 1) b)
    | TSum (a, b) ->
        if Random.State.bool rng then Source.Inl (TSum (a, b), gen_closed (depth - 1) a)
        else Source.Inr (TSum (a, b), gen_closed (depth - 1) b)
    | _ -> gen_closed 0 ty

let sample_terms ~count ~max_depth () =
  let adversarial =
    [
      {
        ty = TBool;
        term =
          Source.Let
            ( "poly_const_false",
              Source.TyLam ("a", Source.Lam ("x", TVar "a", TArrow (TVar "a", TBool), Source.Lam ("y", TVar "a", TBool, Source.Bool false))),
              Source.App
                ( Source.App (Source.TyApp (Source.Var "poly_const_false", TInt), Source.Int 1),
                  Source.Int 1 ) );
      };
      {
        ty = TInt;
        term =
          Source.Let
            ( "x",
              Source.Roll
                ( TMu ("a", TInt),
                  Source.Fix ("loop", TInt, Source.Var "loop") ),
              Source.Int 0 );
      };
      {
        ty = TInt;
        term = Source.Let ("_", Source.Tick ("alloc", Source.Int 0), Source.Int 1);
      };
    ]
  in
  let rec fill acc i attempts =
    if i = count then List.rev acc
    else if attempts > count * 32 then List.rev acc
    else if i < min count 20 then
      let specimen = List.nth adversarial (i mod List.length adversarial) in
      if Typecheck.is_well_typed specimen.ty specimen.term then fill (specimen :: acc) (i + 1) attempts
      else fill acc i (attempts + 1)
    else
      let ty = gen_typ max_depth in
      let specimen = { ty; term = gen_closed max_depth ty } in
      if Typecheck.is_well_typed specimen.ty specimen.term then fill (specimen :: acc) (i + 1) attempts
      else fill acc i (attempts + 1)
  in
  fill [] 0 0

let shrink specimen =
  match specimen.term, specimen.ty with
  | Source.Pair (a, _), ty -> [{ term = a; ty }]
  | Source.If (_, t, e), _ -> [{ specimen with term = t }; { specimen with term = e }]
  | _ -> []

let run_campaign ?(fuel = 256) flags relation ~count ~max_depth () =
  let specimens = sample_terms ~count ~max_depth () in
  let failures, related =
    List.fold_left
      (fun (fails, ok) specimen ->
        let compiled = Pipeline.compile flags specimen.ty specimen.term in
        let cmp = Relation.compare_programs ~fuel relation specimen.ty specimen.term compiled.target_program in
        match cmp.verdict with
        | Relation.Related -> (fails, ok + 1)
        | Relation.Unrelated _ -> ((specimen, cmp) :: fails, ok))
      ([], 0)
      specimens
  in
  { total = List.length specimens; related; failures = List.rev failures }