{-# LANGUAGE ApplicativeDo #-}
module Gyehoek.CPS.Contify
  ( contifyProgram
  ) where

import Control.Monad.Tardis
import Gyehoek.CPS.Syntax
import Gyehoek.Prelude
import qualified Data.HashSet as HS
import Control.Lens.Unsound (adjoin)
import Debug.Pretty.Simple
import qualified Data.HashMap.Strict as H
import Control.Monad.Writer.Lazy
import Control.Monad.Trans.Tardis (liftTardisT)


-- | ain't no way...
-- type T = WriterT (HashSet Name) (Tardis (HashSet Name) (HashSet Name))
type T = TardisT (HashSet Name) (HashSet Name) (Writer (HashSet Name))

evalT :: T a -> a
-- evalT = (`evalTardis` (mempty,mempty)) . fmap fst . runWriterT
evalT = fst . runWriter . (`evalTardisT` (mempty,mempty))

runT :: T a -> (a, HashSet Name)
-- runT = (`evalTardis` (mempty,mempty)) . runWriterT
runT = runWriter . (`evalTardisT` (mempty,mempty))

-- | inline function if it hasn't been used in the past, and won't
-- be used in the future.
tryInline :: Name -> Kappa -> T Kexp
tryInline kname kap = do
  modifyBackwards (HS.insert kname)
  p <- getsPast (HS.member kname)
  modifyForwards (HS.insert kname)
  q <- getsFuture (HS.member kname)
  let c = p || q
  liftTardisT . tell $ if c then HS.singleton kname else mempty
  pure $ if c
    then KexpVar kname
    else KexpKappa kap

getKap :: HashMap Name Abs -> Name -> Maybe Kappa
getKap g kname = g ^? ix kname . #AbsKappa

contify :: HashMap Name Abs -> Exp -> T Exp
contify g = transformM \case
  ExpApply f xs (KexpVar kname) | Just kap <- getKap g kname
   -> ExpApply f xs <$> tryInline kname kap
  ExpPrim p (KexpVar kname) | Just kap <- getKap g kname
   -> ExpPrim p <$> tryInline kname kap
  e -> pure e

contifyProgram :: HoistedProgram -> Eff es HoistedProgram
contifyProgram p = do
  let g = p.bindings
  let (p',contifiedVars) =
        runT $
          traverseOf
            (adjoin
              (#bindings . each . body)
              (#body . body))
            (contify g)
            p
  pTraceShowM contifiedVars
  -- pure $ p' & #bindings %~ H.filterWithKey \k _ -> HS.member k contifiedVars
  pure p'