content: module Grin.Devolve(twiddleGrin,devolveTransform) where
  
  import Control.Monad.Identity
  import Control.Monad.RWS
  import Data.Functor
  import Data.IORef
  import Data.Maybe
  import qualified Data.Map as Map
  import qualified Data.Set as Set
  
  import Grin.Grin
  import Grin.Noodle
  import Options (verbose,fopts)
  import Support.FreeVars
  import Support.Transform
  import Util.Gen
  import Util.SetLike
  import qualified FlagOpts as FO
  
  {-# NOINLINE devolveTransform #-}
  devolveTransform = transformParms {
      transformDumpProgress = verbose,
      transformCategory = "Devolve",
      transformPass = "Grin",
      transformOperation = devolveGrin
      }
  
  -- devolve grin into a form in which it can be readily converted into C code
  -- This lifts any local functions which are ever called in a non-tail-calllike form
  -- to the top level.
  
  devolveGrin :: Grin -> IO Grin
  devolveGrin grin = do
      col <- newIORef []
      let g (n,l :-> r) = f r >>= \r -> return (n,l :-> r)
          f lt@Let { expDefs = defs, expBody = body, .. } = do
              let iterZ :: Bool -> Map.Map Tag (Set.Set Val) -> [FuncDef] -> Map.Map Tag (Set.Set Val)
                  iterZ b pmap (fd@FuncDef { funcDefName = name, funcDefBody = as :-> r }:fs) = iterZ (b || xs' /= xs) (Map.insert name xs pmap) fs where
                      xs = Set.unions $ xs':catMaybes [ Map.lookup t pmap | t <- Set.toList $ freeVars fd]
                      xs' = maybe Set.empty id (Map.lookup name pmap)
                  iterZ True pmap [] = iterZ False pmap defs
                  iterZ False pmap [] = pmap
  
                  nndefs = [ fd | fd <- defs, funcDefName fd `Set.member` expNonNormal ]
                  pmap = iterZ False (fromList [ (funcDefName fd, fromList [ Var x y | (x,y) <- Set.toList $ freeVars (funcDefBody fd), x > v0]) | fd <- nndefs ]) nndefs
  
                  (nmaps,rmaps) = splitEither (map z defs)
                  z fd@FuncDef { funcDefName = name, funcDefBody = as :-> r }
                      | name `Set.member` expNonNormal = Left ((name,(as ++ xs) :-> pr),xs)
                      | otherwise = Right fd { funcDefBody = as :-> pr }
                    where xs = maybe [] Set.toList $ Map.lookup name pmap
                          pr = runIdentity $ proc r
                  proc (App a as t) | Just xs <- Map.lookup a pmap = return (App a (as ++ Set.toList xs) t)
                  proc e = mapExpExp proc e
              --mapM_ print (Map.toList pmap)
              nmaps <- mapM (g . fst) nmaps
              modifyIORef col (++ nmaps)
              updateLetProps <$> mapExpExp f lt { expDefs = rmaps, expBody = runIdentity $ proc body }
          f e = mapExpExp f e
      nf <- mapM g (grinFuncs grin)
      lf <- readIORef col
      let ntenv = extendTyEnv [ createFuncDef False x y | (x,y) <- lf ] (grinTypeEnv grin)
      return $  setGrinFunctions (lf ++ nf) grin { grinPhase = PostDevolve, grinTypeEnv = ntenv }
      --if null lf then return ng else devolveGrin ng
      --if null lf then return ng else devolveGrin ng
  
  -- twiddle does some final clean up before translation to C
  -- it replaces unused arguments with 'v0' and adds GC notations
  
  data Env = Env {
      envMap   :: Map.Map Var Var,
      envRoots :: Set.Set Val,
      envVar   :: Var
      }
  
  newtype R a = R (RWS Env (Set.Set Var) () a)
      deriving(Monad,Functor,MonadReader Env,MonadWriter (Set.Set Var))
  
  runR (R x) = fst $ evalRWS x Env { envRoots = mempty, envMap = mempty, envVar = v1 } ()
  
  class Twiddle a where
      twiddle :: a -> R a
      twiddle a = return a
  
  instance Twiddle Exp where
      twiddle = twiddleExp
  
  instance Twiddle Val where
      twiddle = twiddleVal
  
  instance Twiddle a => Twiddle [a] where
      twiddle xs = mapM twiddle xs
  
  twiddleExp e = f e where
  --    f (BaseOp Promote vs :>>= rest) = f (Return vs :>>= rest)
  --    f (BaseOp Demote vs :>>= rest) = f (Return vs :>>= rest)
      f (x :>>= lam) | fopts FO.Jgc && isAllocing x = do
          roots <- asks envRoots
          let nroots = Set.fromList [ Var v t | (v,t) <- Set.toList (freeVars (if isUsing x then ([] :-> x :>>= lam) else lam)), isNode t, v > v0] Set.\\ roots
          local (\e -> e { envRoots = envRoots e `Set.union` nroots}) $ do
              ne <- return (:>>=) `ap` twiddle x `ap` twiddle lam
              return $ gcRoots (Set.toList nroots) ne
      f (x :>>= lam) = return (:>>=) `ap` twiddle x `ap` twiddle lam
      f l@Let {} = do
          ds <- twiddle (expDefs l)
          b <- twiddle (expBody l)
          roots <- asks envRoots
          let nroots = Set.fromList [ Var v t | (v,t) <- Set.toList (freeVars l), isNode t, v > v0] Set.\\ roots
          return $ gcRoots (Set.toList nroots) (updateLetProps $ l { expDefs = ds, expBody = b })
      f (App a vs t) | fopts FO.Jgc = do
          roots <- asks envRoots
          let nroots = Set.fromList [ Var v t | x <- vs, (v,t) <- Set.toList (freeVars x), isNode t, v > v0] Set.\\ roots
          local (\e -> e { envRoots = envRoots e `Set.union` nroots}) $ do
              ne <- return (App a) `ap` mapM twiddleVal vs `ap` return t
              return $ gcRoots (Set.toList nroots) ne
      f (Case v as) = return Case `ap` twiddle v `ap` twiddle as
      f x | fopts FO.Jgc && isUsing x && isAllocing x = do
          roots <- asks envRoots
          let nroots = Set.fromList [ Var v t | (v,t) <- Set.toList (freeVars x), isNode t, v > v0] Set.\\ roots
          local (\e -> e { envRoots = envRoots e `Set.union` nroots}) $ do
              ne <- mapExpVal twiddleVal x
              return $ gcRoots (Set.toList nroots) ne
      f n = do e <- mapExpVal twiddleVal n ; mapExpExp twiddle e
  
      isUsing (BaseOp StoreNode {} _) = True
      isUsing Alloc {} = True
      isUsing _ = False
  
      isAllocing (BaseOp StoreNode {} _) = True
      isAllocing (BaseOp Eval {} _) = True
      isAllocing (Return [Var {}]) = False
      isAllocing (Return [NodeC {}]) = True
      isAllocing App {} = True
      isAllocing Call {} = True
      isAllocing Let {} = True
      isAllocing (Case _ as) = any isAllocing [ b | _ :-> b <- as]
      isAllocing Alloc {} = True
      isAllocing (e :>>= _ :-> y) = isAllocing e || isAllocing y
      isAllocing _ = False
  
      gcRoots [] x = x
      gcRoots xs e = GcRoots xs e
  
      isNode TyNode = True
      isNode TyINode = True
      isNode (TyPtr TyNode) = True
      isNode (TyPtr TyINode) = True
      isNode _ = False
  
  instance Twiddle Lam where
      twiddle (vs :-> y) = do
          let fvs = freeVars vs
          (y,uv) <- censor (Set.filter (`notElem` fvs)) $ listen (twiddle y)
          let fvp' = Map.fromList $ concatMap (\v -> if v `Set.member` uv then [] else [(v,v0)]) fvs
          vs <- censor (const mempty) . local (\e -> e { envMap = fvp' }) $ twiddle vs
          return (vs :-> y)
  --    twiddle (vs :-> y) = do
  --        cv <- asks envVar
  --        let fvp = Map.fromList $ zip fvs [cv ..]
  --            fvs = freeVars vs
  --        local (\e -> e { envVar = head $ drop (length fvs) [cv .. ], envMap = fvp `Map.union` envMap e }) $ do
  --        (y,uv) <- censor (Set.filter (`notElem` take (length fvs) [cv .. ])) $ listen (twiddle y)
  --        let fvp' = fmap (\v -> if v `Set.member` uv then v else v0) fvp
  --        vs <- censor (const mempty) . local (\e -> e { envMap = fvp' }) $ twiddle vs
  --        return (vs :-> y)
  
  twiddleGrin grin = grinFunctions_s fs' grin where
      fs' = runR . twiddle  $ grinFunctions grin
  
  instance Twiddle FuncDef where
      twiddle = funcDefBody_uM twiddle
  
  twiddleVal x = f x where
      f var@(Var v ty) = do
          em <- asks envMap
          case Map.lookup v em of
              Just n -> tell (Set.singleton n) >> return (Var n ty)
              Nothing -> tell (Set.singleton v) >> return var
      f x = mapValVal f x
  