{-# LANGUAGE OverloadedLists #-}
module Gyehoek.CPS.Close
  ( closeProgram
  ) where

import Gyehoek.CPS.Syntax
import Gyehoek.GenSym
import Gyehoek.Prelude


close :: GenSym :> es => Exp -> Eff es Exp
close = transformM \case
  ExpLetRec [(f, AbsLambda lam@(MkLambda bs kb m))] e -> do
    f_code <- gensym' @Name $ f ^. _Wrapped'. to (<> "-code")
    -- it would probably be most sane to generate a symbol for `env`,
    -- but we're reusing the lambda binding so we don't have to
    -- explicitly substitute recursive calls.
    let frees = freeWithBound' [f] lam
    let m' = ifoldr
             (\n x q -> [cps|(prim (env-ref #{f} #{n})
                                (κ (#{x}) #{q}))|])
             m frees
    pure [cps|
      (letrec ((#{f_code} (λ (#{f} ##{bs} #{kb})
                            #{m'})))
        (prim (make-closure ($ #{f_code}) ##{frees})
          (κ (#{f}) #{e})))
    |]

  ExpApply f xs ktail -> do
    code <- gensym' @Name "code"
    pure [cps|
      (prim (env-code #{f})
        (κ (#{code})
          (#{code} #{f} ##{xs} #{ktail})))
    |]

  e -> pure e

closeProgram :: GenSym :> es => Program -> Eff es Program
closeProgram = traverseOf #body close