open Common
open Anf

(** Tarjan's algorithm for finding strongly connected components in a graph.
    Adapted from the pseudocode found on Wikipedia:
    https://en.wikipedia.org/wiki/Tarjan%27s_strongly_connected_components_algorithm
*)
module Tarjan : sig
  type graph

  val empty : graph
  (** returns an empty graph *)

  val add : int -> int list -> graph -> graph
  (** adds a vertex with its successors to the graph *)

  val of_adjacency_list : (int * int list) list -> graph
  (** converts an adjacency list into a graph *)

  val tarjan : graph -> int list list
  (** calculates the strongly connected components of the graph *)
end = struct
  module IntSet = Set.Make (Int)
  module IntMap = Map.Make (Int)

  type graph = int list IntMap.t

  type state = {
    mutable index : int;
    mutable stack : int list;
    mutable indices : int IntMap.t;
    mutable lowlinks : int IntMap.t;
    mutable on_stack : IntSet.t;
    mutable sccs : int list list;
  }

  let empty : graph = IntMap.empty
  let add : int -> int list -> graph -> graph = IntMap.add

  let of_adjacency_list (graph : (int * int list) list) : graph =
    List.fold_left (fun g (v, vs) -> add v vs g) empty graph

  let create_state () =
    {
      index = 0;
      stack = [];
      indices = IntMap.empty;
      lowlinks = IntMap.empty;
      on_stack = IntSet.empty;
      sccs = [];
    }

  let rec strongconnect state v graph =
    state.indices <- IntMap.add v state.index state.indices;
    state.lowlinks <- IntMap.add v state.index state.lowlinks;
    state.index <- state.index + 1;

    state.stack <- v :: state.stack;
    state.on_stack <- IntSet.add v state.on_stack;

    let successors = IntMap.find_opt v graph |> Option.value ~default:[] in
    List.iter
      (fun w ->
        if not (IntMap.mem w state.indices) then (
          strongconnect state w graph;
          let lowlink_v = IntMap.find v state.lowlinks in
          let lowlink_w = IntMap.find w state.lowlinks in
          state.lowlinks <-
            IntMap.add v (min lowlink_v lowlink_w) state.lowlinks)
        else if IntSet.mem w state.on_stack then
          let lowlink_v = IntMap.find v state.lowlinks in
          let index_w = IntMap.find w state.indices in
          state.lowlinks <- IntMap.add v (min lowlink_v index_w) state.lowlinks)
      successors;

    if IntMap.find v state.lowlinks = IntMap.find v state.indices then
      let rec pop_scc acc =
        match state.stack with
        | [] -> failwith "Unexpected empty stack"
        | x :: xs ->
            state.stack <- xs;
            state.on_stack <- IntSet.remove x state.on_stack;
            let acc = x :: acc in
            if x = v then acc else pop_scc acc
      in
      let scc = pop_scc [] in
      state.sccs <- scc :: state.sccs

  let tarjan graph =
    let state = create_state () in
    IntMap.iter
      (fun v _ ->
        if not (IntMap.mem v state.indices) then strongconnect state v graph)
      graph;
    state.sccs
end

let dependency_graph (bindings : 'a binding list) : Tarjan.graph =
  let vertices = bindings |> List.map (fun b -> b.name) in
  let hash =
    vertices
    |> List.to_seq
    |> Seq.mapi (fun i name -> (name, i))
    |> Hashtbl.of_seq
  in
  let adjacency_list =
    bindings
    |> List.map (fun b -> (b.name, Anf.free_vars_cexpr b.cexpr))
    |> List.map (fun (name, free_vars) ->
           let node = Hashtbl.find hash name in
           let adjacent_nodes =
             List.filter_map (Hashtbl.find_opt hash) free_vars
           in
           (node, adjacent_nodes))
  in
  Tarjan.of_adjacency_list adjacency_list

let calculate_strongly_connected (graph : Tarjan.graph) : int list list =
  Tarjan.tarjan graph

let transform_bindings (bindings : 'a Anf.binding list) (e : 'a Anf.expr) :
    'a Anf.expr =
  bindings
  |> dependency_graph
  |> calculate_strongly_connected
  |> List.map (fun scc -> List.map (fun i -> List.nth bindings i) scc)
  |> List.fold_left (fun e bs -> { e with expr_desc = Letrec (bs, e) }) e

let rec is_lambda (c : 'a Anf.cexpr) : bool =
  match c.cexpr_desc with
  | Lam _ -> true
  | Ascribe (c, _) when is_lambda c -> true
  | AscribeDyn (c, _, _) when is_lambda c -> true
  | _ -> false

let rec purify_expr (e : 'a Anf.expr) : 'a Anf.expr =
  match e.expr_desc with
  | Ret c when is_lambda c ->
      let name = Gensym.fresh "anon" in
      let binding = { name; cexpr = c } in
      let atom = { atom_desc = Var name; atom_meta = e.expr_meta } in
      let cexpr = { cexpr_desc = Atom atom; cexpr_meta = e.expr_meta } in
      let body = { e with expr_desc = Ret cexpr } in
      purify_expr { e with expr_desc = Letrec ([ binding ], body) }
  | Let (b, e') when is_lambda b.cexpr ->
      purify_expr { e with expr_desc = Letrec ([ b ], purify_expr e') }
  | Ret c -> { e with expr_desc = Ret (purify_cexpr c) }
  | Let (b, e') -> { e with expr_desc = Let (purify_binding b, purify_expr e') }
  | Letrec (bs, e) ->
      transform_bindings (List.map purify_binding bs) (purify_expr e)
  | Seq es -> { e with expr_desc = Seq (List.map purify_expr es) }

and purify_cexpr (c : 'a Anf.cexpr) : 'a Anf.cexpr =
  match c.cexpr_desc with
  | Atom a -> { c with cexpr_desc = Atom a }
  | App (a1, atoms) ->
      { c with cexpr_desc = App (purify_atom a1, List.map purify_atom atoms) }
  | Lam (args, e) -> { c with cexpr_desc = Lam (args, purify_expr e) }
  | Prim1 (op, c') -> { c with cexpr_desc = Prim1 (op, purify_cexpr c') }
  | Prim2 (op, a1, a2) ->
      { c with cexpr_desc = Prim2 (op, purify_atom a1, purify_atom a2) }
  | If (a, e1, e2) ->
      { c with cexpr_desc = If (purify_atom a, purify_expr e1, purify_expr e2) }
  | Match (a, pats, ty) ->
      {
        c with
        cexpr_desc = Match (purify_atom a, List.map purify_pattern pats, ty);
      }
  | Tuple atoms -> { c with cexpr_desc = Tuple (List.map purify_atom atoms) }
  | Get (a, i) -> { c with cexpr_desc = Get (purify_cexpr a, i) }
  | Vect (a1, a2) ->
      { c with cexpr_desc = Vect (purify_atom a1, purify_atom a2) }
  | VectGet (a1, a2) ->
      { c with cexpr_desc = VectGet (purify_atom a1, purify_atom a2) }
  | VectSet (a1, a2, a3) ->
      {
        c with
        cexpr_desc = VectSet (purify_atom a1, purify_atom a2, purify_atom a3);
      }
  | VectLength a -> { c with cexpr_desc = VectLength (purify_cexpr a) }
  | Ascribe (a, t) -> { c with cexpr_desc = Ascribe (purify_cexpr a, t) }
  | AscribeDyn (c', a, inv) ->
      { c with cexpr_desc = AscribeDyn (purify_cexpr c', purify_atom a, inv) }
  | For (x, a1, a2, e) ->
      {
        c with
        cexpr_desc = For (x, purify_atom a1, purify_atom a2, purify_expr e);
      }
  | Variant (name, atoms, t) ->
      { c with cexpr_desc = Variant (name, List.map purify_atom atoms, t) }

and purify_atom (a : 'a Anf.atom) : 'a Anf.atom =
  match a.atom_desc with
  | Var x -> { a with atom_desc = Var x }
  | Const c -> { a with atom_desc = Const c }

and purify_binding (b : 'a Anf.binding) : 'a Anf.binding =
  { b with cexpr = purify_cexpr b.cexpr }

and purify_pattern (p : 'a Anf.pattern) : 'a Anf.pattern =
  { p with body = purify_expr p.body }
