-- | Basic data flow analysis over the Haste AST.
module Haste.AST.FlowAnalysis (
    Strict (..), VarInfo (..), InfoMap, ArgMap,
    mkVarInfo, nullInfo, mergeVarInfos, findVarInfos
  ) where
import Haste.AST.Syntax
import Control.Monad.State
import Data.List (foldl', sort, group)
import qualified Data.Set as S
import qualified Data.Map as M

data Strict = Strict | Unknown
  deriving (Eq, Show)

data VarInfo = VarInfo {
    -- | Is this var always strict?
    varStrict :: Strict,

    -- | Does this var always have the same, statically known, arity?
    --   @Nothing@, if var is not a function.
    varArity  :: Maybe Int,

    -- | What are the sources of this var? I.e. for each occurrence where the
    --   var is the LHS of an assignment or substitution (including function
    --   call), add the RHS to this list.
    varAlias  :: [Var]
  } deriving (Eq, Show)

mkVarInfo :: Strict -> Maybe Int -> VarInfo
mkVarInfo s ar = VarInfo s ar []

type InfoMap = M.Map Var VarInfo
type ArgMap = M.Map Var [Var]

nullInfo :: VarInfo
nullInfo = VarInfo Unknown Nothing []

-- | Merge two var infos. The aliases of the infos are simply concatenated
--   and filtered for duplicates. For concrete information, if the two infos
--   disagree on any field, that field becomes unknown.
mergeVarInfos :: VarInfo -> VarInfo -> VarInfo
mergeVarInfos a b = VarInfo strict' arity' (varAlias a `merge` varAlias b)
  where
    strict'
      | varStrict a == varStrict b = varStrict a
      | otherwise                  = Unknown
    arity'
      | varArity a == varArity b   = varArity a
      | otherwise                  = Nothing
    merge as bs = map head . group . sort $ as ++ bs

-- | Merge the infos of a var with that of all of its aliases.
resolveVarInfo :: InfoMap -> Var -> Maybe VarInfo
resolveVarInfo m var = go (S.singleton var) var
  where
    go seen v = do
      nfo <- M.lookup v m
      let xs'   = filter (not . (`S.member` seen)) (varAlias nfo)
          seen' = foldl' (flip S.insert) seen xs'
      nfos <- mapM (go seen') xs'
      case foldl' mergeVarInfos (nfo {varAlias = xs'}) nfos of
        VarInfo Unknown Nothing _ -> Nothing
        nfo'                      -> return nfo'

type EvalM = State InfoMap

-- | Figure out as much statically known information as possible about all
--   vars in a program. The returned map will be collapsed so that recursive
--   lookups are not necessary, and the 'varAlias' field of each info will
--   be empty.
--
--   TODO: track return values, tail call status
findVarInfos :: ArgMap -> Stm -> InfoMap
findVarInfos argmap = mergeAll . snd . flip runState M.empty . goS
  where
    mergeAll m = M.foldWithKey resolve m m

    resolve v _ m =
      case resolveVarInfo m v of
        Just nfo -> M.insert v (nfo {varAlias = []}) m
        _        -> m

    goS Stop = return ()
    goS Cont = return ()
    goS (Forever s) = goS s
    goS (Case c d as next) = do
      goE c
      goS d
      mapM_ (\(e, s) -> goE e >> goS s) as
      goS next
    goS (Assign (NewVar _ v) rhs next)       = handleAssign v rhs >> goS next
    goS (Assign (LhsExp _ (Var v)) rhs next) = handleAssign v rhs >> goS next
    goS (Assign (LhsExp _ lhs) rhs next)     = goE lhs >> goE rhs >> goS next
    goS (Return ex)                          = goE ex
    goS (ThunkRet ex)                        = goE ex
    goS (Tailcall ex)                        = goE ex

    handleAssign v (Var v')        = v `dependsOn` v'
    handleAssign v (Eval (Var v')) = v `dependsOn` v'
    handleAssign v (Fun args body) = do
      v `hasInfo` mkVarInfo Strict (Just $ length args)
      goS body
    handleAssign v (Lit _)         = v `hasInfo` mkVarInfo Strict Nothing
    handleAssign v (JSLit _)       = v `hasInfo` mkVarInfo Strict Nothing
    handleAssign v rhs             = v `hasInfo` nullInfo >> goE rhs

    -- If a function is ever stored anywhere except as a pure alias, we must
    -- immediately stop assuming things about it, since we now have no idea
    -- where it may be called. If it is not stored, we know that all calls to
    -- it will be direct, so it is safe to assume things about its arguments
    -- based on our static analysis.
    -- To be safe - but overly conservative - we set the
    -- info of any arguments of each function aliased by a var that occurs
    -- outside a permitted context (function call or synonymous assignment)
    -- to nullInfo.
    goE (Var v)                  = setArgsOf v nullInfo (S.singleton v)
    goE (Lit _)                  = return ()
    goE (JSLit _)                = return ()
    goE (Not ex)                 = goE ex
    goE (BinOp _ a b)            = goE a >> goE b
    goE (Fun _ body)             = goS body
    goE (Call _ _ (Var f) xs)    = handleCall f xs
    goE (Call _ _ (Fun as b) xs) = handleArgs as xs >> goS b
    goE (Call _ _ f xs)          = mapM_ goE (f:xs)
    goE (Index x ix)             = goE x >> goE ix
    goE (Arr xs)                 = mapM_ goE xs
    goE (Member x _)             = goE x
    goE (Obj xs)                 = mapM_ goE (map snd xs)
    goE (AssignEx (Var v) rhs)   = handleAssign v rhs
    goE (AssignEx lhs rhs)       = goE lhs >> goE rhs
    goE (IfEx c l r)             = mapM_ goE [c, l, r]
    goE (Eval ex)                = goE ex
    goE (Thunk _ body)           = goS body

    handleArgs as xs = sequence_ (zipWith handleAssign as xs)

    handleCall f argexs =
      case M.lookup f argmap of
        Just argvars -> handleArgs argvars argexs
        _            -> return ()

    -- Finds the functions (if any) aliased by v and sets the varinfo of their
    -- arguments to @nfo@.
    setArgsOf v nfo seen = do
      case M.lookup v argmap of
        Just args -> mapM_ (`hasInfo` nullInfo) args
        _         -> do
          nfos <- get
          case M.lookup v nfos of
            Nothing   -> return ()
            Just nfo' -> do
              let aliases = filter (not . (`S.member` seen)) (varAlias nfo')
                  seen' = foldl' (flip S.insert) seen aliases
              mapM_ (\v' -> setArgsOf v' nfo seen') aliases

hasInfo :: Var -> VarInfo -> EvalM ()
hasInfo v nfo = do
    nfos <- get
    put $ M.alter upd v nfos
  where
    upd (Just nfo') = Just $ mergeVarInfos nfo nfo'
    upd _           = Just nfo

dependsOn :: Var -> Var -> EvalM ()
dependsOn v v' = do
    nfos <- get
    put $ M.alter upd v nfos
  where
    upd (Just nfo) = Just $ nfo {varAlias = v' : varAlias nfo}
    upd _          = Just $ nullInfo {varAlias = [v']}