structure CPS =
struct
  fun toCPS (e : Syntax.lexp) (cont : Syntax.value -> Syntax.cexp) : Syntax.cexp =
    case e of
      Syntax.LVar v => cont (Syntax.VVar v)
    | Syntax.LFn (v, expr) =>
        let
          val fnName = GenSym.new ()
          val k = GenSym.new ()
        in
          Syntax.CFix
            ([(fnName, [v, k], toCPS expr (fn ret => Syntax.CApp (Syntax.VVar k, [ret])))],
             cont (Syntax.VVar fnName))
        end
    | Syntax.LFix (decls, body) =>
        Syntax.CFix
          ( map
              (fn (name, arg, expr) =>
                let val w = GenSym.new ()
                in
                  ( name
                  , [arg, w]
                  , toCPS expr (fn z => Syntax.CApp (Syntax.VVar w, [z]))
                  )
                end)
              decls
          , toCPS body cont
          )
    | Syntax.LApp (Syntax.LPrim primop, Syntax.LRecord args) =>
        let
          val temp = GenSym.new ()
          fun go [] acc = Syntax.CPrimop (primop, rev acc, [temp], [cont (Syntax.VVar temp)])
            | go (arg :: args) acc = toCPS arg (fn arg' => go args (arg' :: acc))
        in go args []
        end
    | Syntax.LApp (Syntax.LPrim primop, arg) => toCPS (Syntax.LApp (Syntax.LPrim primop, Syntax.LRecord [arg])) cont
    | Syntax.LApp (f, x) =>
        let
          val addr = GenSym.new ()
          val arg = GenSym.new ()
        in Syntax.CFix
          ([(addr, [arg], cont (Syntax.VVar arg))],
           toCPS f (fn f' =>
             toCPS x (fn x' =>
               Syntax.CApp (f', [x', Syntax.VVar addr]))))
        end
    | Syntax.LInt i => cont (Syntax.VInt i)
    | Syntax.LString s =>
        let val temp = GenSym.new ()
        in
          Syntax.CRecord
            ( [(map (fn c => (Syntax.VInt (Char.ord c), [])) (String.explode s), temp)]
            , cont (Syntax.VVar temp)
            )
        end
    | Syntax.LRecord [] => cont (Syntax.VInt 0)
    | Syntax.LSelect (i, expr) =>
        let val temp = GenSym.new ()
        in toCPS expr (fn x => Syntax.CSelect (i, x, temp, cont (Syntax.VVar temp)))
        end
    | Syntax.LRecord exprs =>
        let
          fun go [] vars =
                let val temp = GenSym.new ()
                in Syntax.CRecord ([(map (fn v => (v, [])) (rev vars), temp)], cont (Syntax.VVar temp))
                end
            | go (expr :: exprs) vars =
                toCPS expr (fn v => go exprs (v :: vars))
        in go exprs []
        end
    | Syntax.LSwitch (expr, arms, otherwise) =>
        let
          val sortedArms = Sort.sort (fn ((x, _), (y, _)) => Int.compare (x, y)) arms
          val contAddr = GenSym.new ()
          val otherwiseAddr = GenSym.new ()
          val arg = GenSym.new ()
          fun go _ [] cont = Syntax.CApp (Syntax.VVar otherwiseAddr, [])
            | go v [(x, arm)] cont =
                (case otherwise of
                  NONE => toCPS arm cont
                | SOME _ =>
                    let val b = GenSym.new ()
                    in
                      Syntax.CPrimop
                        ( Syntax.PEq
                        , [v, Syntax.VInt x]
                        , [b]
                        , [ Syntax.CPrimop
                              ( Syntax.PIf
                              , [Syntax.VVar b]
                              , []
                              , [toCPS arm cont, Syntax.CApp (Syntax.VVar otherwiseAddr, [])]
                              )
                          ]
                        )
                    end)
            | go v arms cont =
                let
                  val b = GenSym.new ()
                  val h = length arms div 2
                  val half1 = List.take (arms, h)
                  val half2 = List.drop (arms, h)
                  val (x, _) = hd half2
                in
                  Syntax.CPrimop
                    ( Syntax.PLess
                    , [v, Syntax.VInt x]
                    , [b]
                    , [ Syntax.CPrimop
                          ( Syntax.PIf
                          , [Syntax.VVar b]
                          , []
                          , [go v half1 cont, go v half2 cont]
                          )
                      ]
                    )
                end
          fun contFunc x = Syntax.CApp (Syntax.VVar contAddr, [x])
          val fixFuncs = [(contAddr, [arg], cont (Syntax.VVar arg))]
          val fixFuncs =
            case otherwise of
              NONE => fixFuncs
            | SOME otherwise => (otherwiseAddr, [], toCPS otherwise contFunc) :: fixFuncs
        in
          Syntax.CFix
            ( fixFuncs
            , toCPS expr (fn v => go v sortedArms contFunc)
            )
        end
    | _ => raise Fail ("malformed expression " ^ Syntax.lexpToString e)

  fun hoist (expr : Syntax.cexp) : Syntax.cexp =
    let
      fun exprs (Syntax.CRecord (records, k)) = Syntax.CRecord (records, exprs k)
        | exprs (Syntax.CSelect (i, arg, res, k)) = Syntax.CSelect (i, arg, res, exprs k)
        | exprs (expr as Syntax.CApp _) = expr
        | exprs (Syntax.CFix (_, body)) = exprs body
        | exprs (Syntax.CPrimop (p, args, res, ks)) = Syntax.CPrimop (p, args, res, (map exprs ks))

      fun funs (Syntax.CRecord (_, k)) acc = funs k acc
        | funs (Syntax.CSelect (_, _, _, k)) acc = funs k acc
        | funs (Syntax.CApp _) acc = acc
        | funs (Syntax.CFix (fs, body)) acc =
            funs body (foldl (fn ((fName, fVars, fBody), acc) => (fName, fVars, exprs fBody) :: funs fBody acc) acc fs)
        | funs (Syntax.CPrimop (_, _, _, ks)) acc = foldl (fn (x, acc) => funs x acc) acc ks
      val entryPoint = GenSym.new ()
    in
      Syntax.CFix ((entryPoint, [], exprs expr) :: funs expr [], Syntax.CApp (Syntax.VLabel entryPoint, []))
    end

  structure VarMap = Map (type k = Syntax.var
                          val cmp = Int.compare)

  fun varSet (l : Syntax.var list) : unit VarMap.map = VarMap.fromList (map (fn x => (x, ())) l)

  fun freeVars (expr : Syntax.cexp) : unit VarMap.map =
    case expr of
      Syntax.CRecord (records, k) =>
        let
          val kFreeVars = freeVars k
          val argFreeVars =
            foldl
              (fn ((args, _), acc) =>
                VarMap.union acc (varSet (List.mapPartial (fn (Syntax.VVar v, _) => SOME v | _ => NONE) args)))
              VarMap.empty
              records
          val results =
            foldl
              (fn ((_, res), acc) =>
                VarMap.insert res () acc)
              VarMap.empty
              records
        in
          VarMap.difference (VarMap.union kFreeVars argFreeVars) results
        end
    | Syntax.CSelect (_, arg, res, k) =>
        let
          val argFreeVars =
            case arg of
              Syntax.VVar v => varSet [v]
            | _ => VarMap.empty
          val kFreeVars = VarMap.delete res (freeVars k)
        in
          VarMap.union argFreeVars kFreeVars
        end
    | Syntax.CApp (func, args) =>
        let
          val funcFreeVars =
            case func of
              Syntax.VVar v => varSet [v]
            | _ => VarMap.empty
          val argFreeVars = varSet (List.mapPartial (fn Syntax.VVar v => SOME v | _ => NONE) args)
        in
          VarMap.union funcFreeVars argFreeVars
        end
    | Syntax.CFix (funs, body) =>
        let
          val names = varSet (map (fn (name, _, _) => name) funs)
          val funsFreeVars = map (fn (name, args, fixBody) => VarMap.difference (freeVars fixBody) (varSet args)) funs
          val bodyFreeVars = freeVars body
        in
          foldl (fn (x, acc) => VarMap.union acc (VarMap.difference x names)) VarMap.empty (bodyFreeVars :: funsFreeVars)
        end
    | Syntax.CPrimop (_, args, res, ks) =>
        let
          val argFreeVars = varSet (List.mapPartial (fn Syntax.VVar v => SOME v | _ => NONE) args)
          val boundVars = varSet res
          val kFreeVars = foldl (fn (x, acc) => VarMap.union acc x) VarMap.empty (map (fn k => VarMap.difference (freeVars k) boundVars) ks)
        in
          VarMap.union argFreeVars kFreeVars
        end

  fun freeVarsClosure (name, args, body) = map (fn (x, _) => x) (VarMap.toList (VarMap.difference (freeVars body) (varSet args)))

  fun enumerate l = ListPair.zip (List.tabulate (length l, (fn x => x)), l)

  fun convertExpr varMap expr =
    let
      fun translate var = getOpt (VarMap.lookup var varMap, var)
      fun translateValue (Syntax.VVar v) = Syntax.VVar (translate v)
        | translateValue v = v
    in
      case expr of
        Syntax.CRecord (records, k) =>
          Syntax.CRecord (map (fn (args, res) => (map (fn (v, p) => (translateValue v, p)) args, res)) records, convertExpr varMap k)
      | Syntax.CSelect (i, arg, res, k) =>
          Syntax.CSelect (i, translateValue arg, res, convertExpr varMap k)
      | Syntax.CApp (func, args) =>
          let
            val temp = GenSym.new ()
            val f = translateValue func
          in Syntax.CSelect (0, f, temp,
               Syntax.CApp (Syntax.VVar temp, f :: map translateValue args))
          end
      | Syntax.CFix (funcs, body) =>
          let
            val convertedFuncs =
              map
                (fn this as (name, args, body) =>
                  let
                    val funcFreeVars = freeVarsClosure this
                    val varMap' = VarMap.union varMap (VarMap.fromList (map (fn v => (v, GenSym.new ())) funcFreeVars))
                    val closure = GenSym.new ()
                    val newBody =
                      foldl
                        (fn ((i, x), acc) => Syntax.CSelect (i + 1, Syntax.VVar closure, valOf (VarMap.lookup x varMap'), acc))
                        (convertExpr varMap' body)
                        (enumerate funcFreeVars)
                  in
                    (GenSym.new (), closure :: args, newBody)
                  end)
                funcs
            val newBody =
              Syntax.CRecord
                ( map
                    (fn (old as (oldName, args, body), (newName, _, _)) =>
                      let val funcFreeVars = freeVarsClosure old
                      in ((Syntax.VLabel newName, []) :: map (fn v => (Syntax.VVar (translate v), [])) funcFreeVars, oldName)
                      end)
                    (ListPair.zip (funcs, convertedFuncs))
                , convertExpr varMap body
                )
          in
            Syntax.CFix (convertedFuncs, newBody)
          end
      | Syntax.CPrimop (p, args, res, ks) =>
          Syntax.CPrimop (p, map translateValue args, res, map (convertExpr varMap) ks)
    end

  fun convertClosures (expr : Syntax.cexp) : Syntax.cexp = hoist (convertExpr VarMap.empty expr)
end