(* Histogram of Tast expr_kind constructors over the reachable program. A
   measurement, not a backend: it answers "what would a whole-program x86
   build actually have to lower for this input", which is the question that
   decides whether whole-program coverage is reachable at all. *)
let tbl : (string, int) Hashtbl.t = Hashtbl.create 64

let bump k =
  Hashtbl.replace tbl k (1 + (try Hashtbl.find tbl k with Not_found -> 0))

let name (k : Flan.Tast.expr_kind) =
  match k with
  | Int _ -> "Int" | Float _ -> "Float" | Bool _ -> "Bool" | Str _ -> "Str"
  | Unit -> "Unit" | Zero _ -> "Zero" | Uninit _ -> "Uninit"
  | Local _ -> "Local" | Global _ -> "Global" | Prim _ -> "Prim"
  | Call _ -> "Call" | FnAddr _ -> "FnAddr" | CallPtr _ -> "CallPtr"
  | Do _ -> "Do" | Let _ -> "Let" | If _ -> "If" | While _ -> "While"
  | Return _ -> "Return" | Break _ -> "Break" | Continue _ -> "Continue"
  | Set _ -> "Set" | Field _ -> "Field" | Addr _ -> "Addr" | Deref _ -> "Deref"
  | Make _ -> "Make" | MakeCase _ -> "MakeCase" | CaseField _ -> "CaseField"
  | Arr _ -> "Arr" | Some_ _ -> "Some" | None_ -> "None" | Match _ -> "Match"
  | UnwrapSome _ -> "UnwrapSome" | Signal _ -> "Signal" | Handled _ -> "Handled"
  | RestartCase _ -> "RestartCase" | WithAlloc _ -> "WithAlloc"
  | InvokeRestart _ -> "InvokeRestart"

let pname (p : Flan.Tast.prim) =
  match p with
  | Add -> "Add" | Sub -> "Sub" | Mul -> "Mul" | Div -> "Div" | Rem -> "Rem"
  | Eq -> "Eq" | Ne -> "Ne" | Lt -> "Lt" | Le -> "Le" | Gt -> "Gt" | Ge -> "Ge"
  | Not -> "Not" | BitAnd -> "BitAnd" | BitOr -> "BitOr" | BitXor -> "BitXor"
  | Shl -> "Shl" | Shr -> "Shr" | Len -> "Len" | At -> "At" | Slice -> "Slice"
  | Bytes -> "Bytes" | BytesToF64 -> "BytesToF64" | BytesToI64 -> "BytesToI64"
  | F64ToBytes -> "F64ToBytes" | I64ToBytes -> "I64ToBytes"
  | StrOfBytes -> "StrOfBytes" | U64ToBytes -> "U64ToBytes"
  | EscapeBytes -> "EscapeBytes" | WriteStdout -> "WriteStdout" | Exit -> "Exit"
  | Argv -> "Argv" | Rt s -> "Rt:" ^ s | SizeOf _ -> "SizeOf"
  | AlignOf _ -> "AlignOf" | AddrOf -> "AddrOf" | Cast _ -> "Cast"

let rec ex (e : Flan.Tast.expr) =
  bump (name e.e);
  match e.e with
  | Prim (p, xs) -> bump ("prim/" ^ pname p); List.iter ex xs
  | Call (_, xs) | Arr xs -> List.iter ex xs
  | Make (_, xs) | MakeCase (_, _, xs) -> List.iter ex xs
  | CallPtr (f, xs) -> ex f; List.iter ex xs
  | Do xs | Handled (_, xs) -> List.iter ex xs
  | Let (bs, body) -> List.iter (fun (_, x) -> ex x) bs; List.iter ex body
  | If (a, b, c) -> ex a; ex b; ex c
  | While (c, b, l) -> ex c; List.iter ex b; List.iter ex l
  | Return (Some x) | Some_ x | Deref x | UnwrapSome x | Field (x, _)
  | CaseField (x, _, _) | Signal (_, _, x) -> ex x
  | Set (p, x) -> pl p; ex x
  | Addr p -> pl p
  | Match (x, arms) ->
    ex x;
    List.iter (fun (a : Flan.Tast.arm) -> List.iter ex a.abody) arms
  | RestartCase (cs, x) ->
    List.iter (fun (c : Flan.Tast.rclause) -> List.iter ex c.rbody) cs; ex x
  | WithAlloc (a, b) -> ex a; List.iter ex b
  | InvokeRestart (_, _, xs, _, _, _) -> List.iter ex xs
  | _ -> ()

and pl (p : Flan.Tast.place) =
  match p with
  | Plocal _ -> bump "place/Plocal"
  | Pglobal _ -> bump "place/Pglobal"
  | Pfield (x, _) -> bump "place/Pfield"; ex x
  | Pindex (x, ys) -> bump "place/Pindex"; ex x; List.iter ex ys
  | Pderef x -> bump "place/Pderef"; ex x

let () =
  let src = Sys.argv.(1) in
  let l =
    Flan.Load.program ~file:src
      (Flan.Parse.program_all (Flan.Reader.read_file src))
  in
  let p = Flan.Check.program_all l.Flan.Load.decls in
  let p, _, _ = Flan.Reach.link l p in
  List.iter
    (fun (f : Flan.Tast.fn) -> List.iter ex f.body; List.iter ex f.fdefers)
    p.Flan.Tast.fns;
  List.iter (fun (g : Flan.Tast.global) -> ex g.Flan.Tast.ginit)
    p.Flan.Tast.globals;
  Printf.printf "%s: %d reachable fns\n" (Filename.basename src)
    (List.length p.Flan.Tast.fns);
  let rows = Hashtbl.fold (fun k v a -> (k, v) :: a) tbl [] in
  let rows = List.sort (fun (a, _) (b, _) -> compare a b) rows in
  List.iter (fun (k, v) -> Printf.printf "  %-24s %d\n" k v) rows