content: {-# LANGUAGE PatternGuards, TupleSections, DoAndIfThenElse, OverloadedStrings #-}
  -- | Optimizations over the JSTarget AST.
  module Haste.AST.Optimize (
      optimizeFun, tryTernary, topLevelInline
    ) where
  import Haste.AST.Syntax
  import Haste.AST.Op
  import Haste.AST.FlowAnalysis
  import Haste.AST.Traversal
  import Haste.Config
  import Control.Applicative
  import Control.Monad
  import qualified Data.Map as M
  import qualified Data.Set as S
  import qualified Data.ByteString.Char8 as BS
  
  -- | Turn tail recursion into loops.
  fixTailCalls :: Var -> Exp -> TravM Exp
  fixTailCalls fun ast = do
      ast' <- assignToSubst ast >>= tailLoopify fun
      mapJS (const True) pure loopify ast'
    where
      loopify (Assign lhs@(NewVar _ f) body next) =
        Assign lhs <$> (assignToSubst body >>= tailLoopify f) <*> pure next
      loopify stm =
        return stm
  
  -- | Apply the given optimization when enabled by the config.
  optWhen :: Bool -> (a -> TravM a) -> a -> TravM a
  optWhen doit opt
    | doit      = opt
    | otherwise = pure
  
  -- TODO: tryTernary may inline calls that would otherwise be in tail position
  --       which is something we'd really like to avoid.
  optimizeFun :: Config -> Var -> Exp -> Exp
  optimizeFun c f ast =
    runTravM $ do
      shrinkCase ast
      >>= inlineAssigns []
      >>= optimizeArrays
      >>= inlineReturns
      >>= zapJSStringConversions
      >>= optimizeThunks
      >>= optimizeArrays
      >>= optWhen (enableTailLoops c) (fixTailCalls f)
      >>= inlineShortJumpTailcall
      >>= optWhen (enableProperTailcalls c) trampoline
      >>= ifReturnToTernary
      >>= smallStepInline
      >>= optWhen (inlineJSPrim c) inlineJSPrimitives
      >>= mapJS (const True) pure (pure . removeNonsenseAssigns)
  
  topLevelInline :: Config -> [Name] -> Stm -> Stm
  topLevelInline c spt ast =
    runTravM $ do
      pure ast
      >>= unevalLits
      >>= inlineIntoEval
      >>= inlineAssigns spt'
      >>= optimizeArrays
      >>= optimizeThunks
      >>= smallStepInline
      >>= optimizeArrays
      >>= zapJSStringConversions
      >>= optWhen (inlineJSPrim c) inlineJSPrimitives
      >>= removeUselessEvals
      >>= inlineAssigns spt'
      >>= smallStepInline
      >>= inlineAssigns spt'
      >>= optWhen (detrampolineThreshold c > 0)
                  (unTrampoline $ detrampolineThreshold c)
      >>= inlineLiteralFunCalls
      >>= optWhen (flowAnalysisOpts c) flowOpts
    where
      spt' = map (\n -> Internal n "" False) spt
  
  
  -- | Attempt to turn two case branches into a ternary operator expression.
  tryTernary :: Var
             -> Exp
             -> Exp
             -> (Stm -> Stm)
             -> [(Exp, Stm -> Stm)]
             -> Maybe Exp
  tryTernary self scrut retEx def [(m, alt)] =
      runTravM opt
    where
      selfOccurs (Exp (Var v) _) = v == self
      selfOccurs _               = False
      def' = def $ Return retEx
      alt' = alt $ Return retEx
      opt = do
        -- Make sure the return expression is used somewhere, then cut away all
        -- useless assignments. If what's left is a single Return statement,
        -- we have a pure expression suitable for use with ?:.
        def'' <- inlineAssigns [] def'
        alt'' <- inlineAssigns [] alt'
        -- If self occurs in either branch, we can't inline or we risk ruining
        -- tail call elimination.
        selfInDef <- occurrences (const True) selfOccurs def''
        selfInAlt <- occurrences (const True) selfOccurs alt''
        case (selfInDef + selfInAlt, def'', alt'') of
          (Never, Return el, Return th) ->
            return $ Just $ IfEx (BinOp Eq scrut m) th el
          _ ->
            return Nothing
  tryTernary _ _ _ _ _ =
    Nothing
  
  -- | Remove bogus assignments of the form @literal = exp@ and @_ = _@,
  --   which may arise from other optimizations.
  removeNonsenseAssigns :: Stm -> Stm
  removeNonsenseAssigns (Assign (LhsExp _ (Lit _)) _ next)          = next
  removeNonsenseAssigns (Assign (LhsExp _ a) b next)       | a == b = next
  removeNonsenseAssigns (Assign (NewVar _ a) (Var b) next) | a == b = next
  removeNonsenseAssigns stm                                         = stm
  
  -- | How many times does an expression satisfying the given predicate occur in
  --   an AST (including jumps)?
  occurrences :: JSTrav ast
              => (ASTNode -> Bool)
              -> (ASTNode -> Bool)
              -> ast
              -> TravM Occs
  occurrences tr p ast =
      foldJS trav count Never ast
    where
      trav n node = tr node && n < Lots -- Stop traversal if we're already >1.
      count n node | p node = pure $ n + Once
      count n _             = pure n
  
  -- | Inline assignments where the assignee is only ever used once.
  --   Does not inline anything into a shared code path, as that would break
  --   things horribly. Also doesn't inline anything from the SPT.
  --   The SPT is only necessary when performing top-level inlining. For anything
  --   else an empty list should be used, as any SPT entries are inevitably going.
  --   to be free vars and thus not inlineable.
  --
  --   Ignores LhsExp assignments, since we only introduce those when we actually
  --   care about the assignment side effect.
  inlineAssigns :: JSTrav ast => [Var] -> ast -> TravM ast
  inlineAssigns spt ast = do
      inlinable <- countVarOccs ast
      mapJS (const True) return (inl inlinable) ast
    where
      -- Assume that an assignment that is never used is actually there for a
      -- good reason and thus not inlinable.
      isSafe i v
        | v `elem` spt              = False
        | Just Once <- M.lookup v i = True
        | otherwise                 = False
  
      atLeastOnce i v
        | Just occs <- M.lookup v i = occs > Never
        | otherwise                 = False
  
      isJSLit (Exp (JSLit {}) _) = True
      isJSLit _                  = False
  
      isVar v (Exp (Var v') _) = v == v'
      isVar _ _                = False
  
      isWritten v (Exp (AssignEx (Var v') _) _)            = v == v'
      isWritten v (Stm (Assign (LhsExp _ (Var v')) _ _) _) = v == v'
      isWritten _ _                                        = False
  
      -- It's OK to inline lambda-likes into anything that's not a loop or a
      -- function (thunks are OK because they only get evaluated once).
      goodInlineTgt = not <$> isLoop .|. isFun
  
      -- Only inline safe-to-reorder vars - if a var is ever assigned to, it is
      -- insanely unsafe to inline it.
      inl safe keep@(Assign (NewVar True v) rhs next) = do
        written <- occurrences (const True) (isWritten v) next
        case rhs of
          _ | written /= Never -> do
            return keep
          -- String lits: inline if used exactly once to avoid blowing up code size
          l@(Lit (LStr _)) | isSafe safe v -> do
            replaceEx (const True) (Var v) l next
  
          -- Other lits: inline if used at least once
          l@(Lit _) | atLeastOnce safe v -> do
            occs <- occurrences (const True) (isVar v) next
            if occs > Never
              then replaceEx (const True) (Var v) l next
              else return keep
  
          -- Variables: inline if used at least once
          v'@(Var {}) | atLeastOnce safe v -> do
            occs <- occurrences (const True) (isVar v) next
            if occs > Never
              then replaceEx (const True) (Var v) v' next
              else return keep
  
          -- Everything else: inline if only used once and doesn't contain
          -- lambda-likes. Lambda-likes may *never* be inlined as they may be
          -- referred to in a *previous* binding and thus be recursive without
          -- us knowing. We could fix this if we collected information about
          -- lambda dependencies properly before starting any inlining,
          -- if the lambda-like occurs exactly once, and this occurrence is not
          -- in a lambda or a loop.but for now we disallow it.
          ex | isSafe safe v -> do
            lambdaoccs <- occurrences (const True) (isLambda .|. isJSLit) ex
            recursive <- occurrences (const True) (isVar v) ex
            case (lambdaoccs, recursive) of
              (Never, Never) -> do
                (reps, next') <- replaceExWithCount (const True) (Var v) ex next
                if reps == 1
                  then return next'
                  else return keep
              (Once, Never)  -> do
                -- If we only made a single replacement in a good inline target,
                -- then we know that everything is OK. If not, we accidentally
                -- inlined into a bad code block, so we discard the result.
                (reps, next') <- replaceExWithCount goodInlineTgt (Var v) ex next
                if reps == 1
                  then return next'
                  else return keep
              _              -> do
                pure keep
          _ -> do
            pure keep
      inl _ keep = do
        pure keep
  
  -- | Remove an occurrence of @ex = E(ex)@ or @lit = ex@.
  --   Only call this for @ex@ which are guaranteed to never be thunks.
  removeUpdate :: (Var -> Bool) -> Stm -> Stm
  removeUpdate p stm@(Assign _ _ next) | isEvalUpd p stm = next
  removeUpdate _ stm                                     = stm
  
  -- | Inline function calls of the form @return (function(as){body})(xs)@.
  --   TODO: 'inlineShortJumpTailcall' performs the same functionality for
  --   named local tail calls - merge?
  inlineLiteralFunCalls :: JSTrav ast => ast -> TravM ast
  inlineLiteralFunCalls ast = do
      mapJS (const True) pure (pure . inline) ast
    where
      inline s@(Return (Call _ (Fast False) (Fun as body) xs)) =
        maybe s id (zipAssign (map (NewVar True) as) xs body)
      inline s@(ThunkRet (Call _ (Fast False) (Fun as body) xs)) =
        maybe s id (zipAssign (map (NewVar True) as) xs body)
      inline s = s
  
  -- | Turn the common pattern @var x = e ; x = E(x)@ into @var x = E(e)@.
  --   should run *after* 'unevalLits'.
  inlineIntoEval :: JSTrav ast => ast -> TravM ast
  inlineIntoEval ast = do
      mapJS (const True) pure (pure . inline) ast
    where
      inline (Assign l@(NewVar _ v) r s@(Assign _ _ next))
        | isEvalUpd (== v) s = Assign l (Eval r) next
      inline stm             = stm
  
  isEvalUpd :: (Var -> Bool) -> Stm -> Bool
  isEvalUpd p (Assign (LhsExp _ (Var v)) (Eval (Var v')) _) = p v && v == v'
  isEvalUpd _ _                                             = False
  
  -- | Turn if(foo) {return bar;} else {return baz;} into return foo ? bar : baz.
  ifReturnToTernary :: JSTrav ast => ast -> TravM ast
  ifReturnToTernary ast = do
      mapJS (const True) return opt ast
    where
      opt (Case cond (Return el) [(ex, Return th)] _) =
        pure $ Return $ IfEx (BinOp Eq cond ex) th el
      opt stm =
        pure stm
  
  -- | Turn occurrences of @[a,b][1]@ or @{... x:b ...}.b@ into b.
  optimizeArrays :: JSTrav ast => ast -> TravM ast
  optimizeArrays ast =
      mapJS (const True) inlEx return ast
    where
      inlEx (Index (Arr xs) (Lit (LNum n))) =
        return $ xs !! truncate n
      inlEx (Member (Obj xs) key) =
        case lookup key xs of
          Just x -> return x
          _      -> error $  "Bad object literal deref: member " ++ show key
                          ++ " doesn't exist!"
      inlEx x =
        return x
  
  -- | Remove calls to E() where we know for sure the argument is already
  --   evaluated.
  removeUselessEvals :: JSTrav ast => ast -> TravM ast
  removeUselessEvals ast = do
      mapJS (const True) return opt ast
    where
      opt (Assign lhs@(NewVar _ l) rhs next) = do
        Assign lhs rhs <$> case rhs of
          Eval r@(Var _) -> removeEval (Var l) next >>= removeEval r
          _ | isHNF rhs  -> removeEval (Var l) next
            | otherwise  -> return next
      opt stm = do
        return stm
  
      isHNF (Lit {})   = True
      isHNF (JSLit {}) = True
      isHNF (Not {})   = True
      isHNF (BinOp {}) = True
      isHNF (Fun {})   = True
      isHNF (Arr {})   = True
      isHNF (Eval {})  = True
      isHNF _          = False
  
      removeEval x = replaceEx (const True) (Eval x) x
  
  -- | Turn toJSStr(unCStr(x)) into x, since rewrite rules absolutely refuse
  --   to work with unpackCString#.
  --   Also turn T(unCStr(x)) into unCStr(x) whenever x is a literal, since
  --   unCStr is evaluated lazily anyway.
  zapJSStringConversions :: JSTrav ast => ast -> TravM ast
  zapJSStringConversions ast =
      mapJS (const True) opt return ast
    where
      opt (Call _ _ (Var (Foreign "toJSStr")) [
             Call _ _ (Var (Foreign "unCStr")) [x]]) =
        return x
      opt (Call _ _ (Var (Foreign "toJSStr")) [
             Eval (Call _ _ (Var (Foreign "unCStr")) [x])]) =
        return x
      opt (Thunk _ (Return x@(Call _ _ (Var (Foreign "unCStr")) [Lit _]))) =
        return x
      opt x =
        return x
  
  -- | Optimize thunks in the following ways:
  --   1. A(thunk(return f), xs)
  --        => A(f, xs)
  --   2. thunk(x@(JSLit s)) | s is a JS function object or marked eager
  --        => x
  --   3. thunk(x@(Lit _))
  --        => x
  --   4. E(thunk(return x))
  --        => x
  --   5. E(x) | x is guaranteed to not be a thunk
  --        => x
  --
  --   Note that #2 depends on the invariant of 'JSLit': a JS literal must not
  --   perform side effects or significant computation.
  optimizeThunks :: JSTrav ast => ast -> TravM ast
  optimizeThunks ast =
      mapJS (const True) optEx return ast
    where
      optEx (Eval x)
        | Just x' <- fromThunkEx x           = return x'
        | definitelyNotThunk x               = return x
      optEx ex@(Thunk _ _)
        | Just l@(JSLit s) <- fromThunkEx ex =
          case maybeExtractStrict s of
            Just s'           -> return $ JSLit s'
            _ | isJSFunDecl s -> return l
              | otherwise     -> return ex
      optEx ex@(Thunk _ _)
        | Just l@(Lit _) <- fromThunkEx ex   = return l
      optEx (Call arity calltype f as)
        | Just f' <- fromThunkEx f           = return $ Call arity calltype f' as
      optEx ex                               = return ex
  
  maybeExtractStrict :: BS.ByteString -> Maybe BS.ByteString
  maybeExtractStrict js
    | "__strict(" `BS.isPrefixOf` js && ")" `BS.isSuffixOf` js =
      Just $ BS.init $ BS.drop 9 js
    | otherwise =
      Nothing
  
  -- | Conservatively approximate whether a given JS literal is a function
  --   declaration or not.
  --
  --   TODO: proper parsing here.
  isJSFunDecl :: BS.ByteString -> Bool
  isJSFunDecl s
    | "function(" `BS.isPrefixOf` s && "}" `BS.isSuffixOf` s  = True
    | "(function(" `BS.isPrefixOf` s && ")" `BS.isSuffixOf` s = True
    | otherwise                                               = False
  
  -- | Unpack the given expression if it's a thunk.
  fromThunk :: Exp -> Maybe Stm
  fromThunk (Thunk _ body) = Just body
  fromThunk _              = Nothing
  
  -- | Unpack the given expression if it's a thunk without internal bindings.
  fromThunkEx :: Exp -> Maybe Exp
  fromThunkEx ex =
    case fromThunk ex of
      Just (Return ex')   -> Just ex'
      Just (ThunkRet ex') -> Just ex'
      _                   -> Nothing
  
  -- | Count the occurrences of all variables in an AST.
  countVarOccs :: JSTrav ast => ast -> TravM (M.Map Var Occs)
  countVarOccs ast = do
      foldJS (\_ _->True) countOccs (M.empty) ast
    where
      updVar (Just occs) = Just (occs+Once)
      updVar _           = Just Once
      updVarAss (Just o) = Just o
      updVarAss _        = Just Never
  
      {-# INLINE countOccs #-}
      countOccs m (Exp (Var v@(Internal _ _ _)) _) =
        pure (M.alter updVar v m)
      countOccs m (Stm (Assign (NewVar _ v) _ _) _) =
        pure (M.alter updVarAss v m)
      countOccs m (Stm (Assign (LhsExp True (Var v)) _ _) _) =
        pure (M.alter updVarAss v m)
      countOccs m _ =
        pure m
  
  -- | May the given expression ever tailcall?
  --   TODO:
  --     Be slightly smarter about handling locally defined functions; always
  --     counting a tailcall from a local as a tailcall from the containing
  --     function seems a bit too restrictive. On the other hand, this makes
  --     only a very slight difference in the number of unnecessary tailcalls
  --     eliminated.
  mayTailcall :: JSTrav ast => ast -> TravM Bool
  mayTailcall ast = do
      foldJS enter countTCs False ast
    where
      enter True _                = False
      enter _ (Exp (Thunk _ _) _) = False
      enter _ (Exp (Fun _ _) _)   = False
      enter _ _                   = True
      countTCs _ (Stm (Tailcall _) _) = return True
      countTCs acc _                  = return acc
  
  -- | Gather a map of all symbols which we know will never make tail calls.
  --   All calls to functions in this set can then safely be de-trampolined.
  gatherNonTailcalling :: Stm -> TravM (S.Set Var)
  gatherNonTailcalling stm = do
      foldJS (\_ _ -> True) countTCs S.empty stm
    where
      countTCs s (Exp (Var v@(Foreign _)) _) = do
        return $ S.insert v s
      countTCs s (Stm (Assign (NewVar _ v) (JSLit {}) _) _) = do
        return $ S.insert v s
      countTCs s (Stm (Assign (NewVar _ v) (Fun _ body) _) _) = do
        tc <- mayTailcall body
        return $ if not tc then S.insert v s else s      
      countTCs s (Stm (Assign (LhsExp True (Var v)) (Fun _ body) _) _) = do
        tc <- mayTailcall body
        return $ if not tc then S.insert v s else s
      countTCs s _ = do
        return s
  
  -- | Remove trampolines wherever possible.
  --   The trampoline machinery has some overhead; two extra activation records
  --   on the stack for a single, non-tailcalling function, to be precise.
  --   We observe that bouncing a function that is guaranteed to never tailcall
  --   is a waste of resources, so we can remove those bounces.
  --   Additionally, tailcalling a function which is guaranteed to not tailcall
  --   in turn is wasteful (see above comment about overhead), so we can
  --   eliminate any such function.
  --   Since the tailcalling machinery grows the stack by a total of three
  --   activation records for an arbitrary string of tailcalling functions,
  --   we can apply this procedure recursively three times and still be
  --   guaranteed to use no more stack frames than we would have without this
  --   optimization.
  --
  --   Note that we can only do this for fast calls, as we need to be able to
  --   know exactly what function is actually called upon saturated application.
  unTrampoline :: Int -> Stm -> TravM Stm
  unTrampoline 0 s = pure s
  unTrampoline n s = do
      ntcs <- gatherNonTailcalling s
      mapJS (const True) (unTr ntcs) (unTC ntcs) s >>= unTrampoline (n-1)
    where
      unTr ntcs (Call ar (Fast True) f@(Var v) xs)
        | v `S.member` ntcs =
          return $ Call ar (Fast False) f xs
      unTr _ c@(Call ar (Fast True) f@(Fun _ body) xs) = do
          tc <- mayTailcall body
          return $ if tc then c else Call ar (Fast False) f xs
      unTr _ x =
          return x
  
      -- If we know for certain that the function we're tailcalling will not
      -- tailcall in turn we should not tailcall it, since that would mean two
      -- activation records on the stack - one for the trampoline and one for
      -- the function itself.
      unTC ntcs (Tailcall c@(Call _ (Fast _) (Var v) _))
        | v `S.member` ntcs =
          return $ Return c
      unTC _ tc@(Tailcall c@(Call _ (Fast _) (Fun _ body) _)) = do
          maytc <- mayTailcall body
          if not maytc then return (Return c) else return tc
      unTC _ x =
          return x
  
  -- | Turn sequences like `v0 = foo; v1 = v0; v2 = v1; return v2;` into a
  --   straightforward `return foo;`.
  --   Ignores LhsExp assignments, since we only introduce those when we actually
  --   care about the assignment side effect.
  inlineReturns :: JSTrav ast => ast -> TravM ast
  inlineReturns ast = do
      mapJS (const True) inl pure ast
    where
      inl (Fun as body)    = Fun as <$> go Nothing body
      inl (Thunk upd body) = Thunk upd <$> go Nothing body
      inl ex               = pure ex
   
      goAlt outside (ex, stm) = (ex,) <$> go outside stm
  
      go outside (Case c d as next) = do
        next' <- go outside next
        case returnLike next' of
          outside'@(Just (Var _, _)) ->
            Case c <$> go outside' d <*> mapM (goAlt outside') as <*> pure Stop
          _ ->
            Case c <$> go Nothing d <*> mapM (goAlt Nothing) as <*> pure next'
      go _ (Forever s) = do
        Forever <$> go Nothing s
      go outside (Assign l@(NewVar _ lhs) r next) = do
        next' <- go outside next
        case (next', outside) of
          (Stop, Just (Var v, ret))   | v == lhs -> return $ ret r
          (Return (Var v), _)         | v == lhs -> return $ Return r
          (ThunkRet (Var v), _)       | v == lhs -> return $ ThunkRet r
          (Assign ll (Var v) Stop, _) | v == lhs -> return $ Assign ll r Stop
          _                                      -> return $ Assign l r next'
      go outside (Assign l r next) = do
        Assign l r <$> go outside next
      go _ stm = do
        return stm
  
  -- | Extract the expression returned from a Return of ThunkRet, as well as
  --   a function to recreate that type of return.
  returnLike :: Stm -> Maybe (Exp, Exp -> Stm)
  returnLike (Return e)   = Just (e, Return)
  returnLike (ThunkRet e) = Just (e, ThunkRet)
  returnLike _            = Nothing
  
  -- | Shrink case statements as much as possible.
  shrinkCase :: JSTrav ast => ast -> TravM ast
  shrinkCase =
      mapJS (const True) pure shrink
    where
      shrink (Case _ def [] next)
        | def == Stop = return next
        | otherwise   = replaceFinalStm next (== Stop) def
      shrink stm      = return stm
  
  -- | Turn any calls in tail position into tailcalls.
  --   Must run after @tailLoopify@ or we won't get loops for simple tail
  --   recursive functions.
  trampoline :: Exp -> TravM Exp
  trampoline = mapJS (pure True) pure bounce
    where
      bounce (Return (Call arity call f args)) = do
        return $ Tailcall $ Call arity call' f args
        where
          call' =
            case call of
              Normal _ -> Normal False
              Fast _   -> Fast False
              c        -> c
      bounce s = do
        return s
  
  -- | Turn tail recursion on the given var into a loop, if possible.
  --   Tail recursive functions that create closures turn into:
  --   function f(a', b', c') {
  --     while(1) {
  --       var r = (function(a, b, c) {
  --         a' = a; b' = b; c' = c;
  --       })(a', b', c');
  --       if(r != null) {
  --         return r;
  --       }
  --     }
  --   }
  tailLoopify :: Var -> Exp -> TravM Exp
  tailLoopify f fun@(Fun args body) = do
      tailrecs <- occurrences (not <$> isLambda) isTailRec body
      if tailrecs > Never
        then do
          needToCopy <- createsClosures body
          case needToCopy of
            True -> do
              let args' = map newName args
                  ret = Return continueSentinel
              b <- mapJS (not <$> isLambda) pure (replaceByAssign ret args') body
              tailcalls <- occurrences (const True) isHiddenTailcall body
              let nn = newName f
                  nv = NewVar False nn
                  body' =
                    Forever $
                    Assign nv (Call 0 (Fast (tailcalls > Never)) (Fun args b)
                                                                 (map Var args')) $
                    Case (Var nn) (Return (Var nn)) [(continueSentinel, Stop)] $
                    Stop
              return $ Fun args' body'
            False -> do
              let c = Cont
              body' <- mapJS (not <$> isLambda) pure (replaceByAssign c args) body
              return $ Fun args (Forever body')
        else do
          return fun
    where
      continueSentinel = Var $ Foreign "__continue"
      isTailRec (Stm (Return (Call _ _ (Var f') _)) _) = f == f'
      isTailRec _                                      = False
  
      isHiddenTailcall (Stm (Return (Call {})) _) = True
      isHiddenTailcall _                          = False
      
      -- Only traverse until we find a closure
      createsClosures = foldJS (\acc _ -> not acc) isClosure False
      isClosure _ (Exp (Fun _ _) _)   = pure True
      isClosure _ (Exp (Thunk _ _) _) = pure True
      isClosure acc _                 = pure acc
  
      -- Assign any changed vars, then loop.
      replaceByAssign end as (Return (Call _ _ (Var f') as')) | f == f' = do
        let (first, second) = foldr assignUnlessEqual (id, end) (zip as as')
        return $ first second
      replaceByAssign _ _ stm =
        return stm
  
      -- Assign an expression to a variable, unless that expression happens to
      -- be the variable itself.
      assignUnlessEqual (v, (Var v')) (next, final)
        | v == v' =
          (next, final)
      assignUnlessEqual (v, x) (next, final)
        | any (x `contains`) args =
          (Assign (NewVar False (newName v)) x . next,
           Assign (LhsExp False (Var v)) (Var $ newName v) final)
        | otherwise =
          (next, Assign (LhsExp False (Var v)) x final)
      
      newName (Internal (Name n mmod) _ _) =
        Internal (Name (BS.cons ' ' n) mmod) "" True
      newName n =
        n
      
      contains (Var v) var          = v == var
      contains (Lit _) _            = False
      contains (JSLit _) _          = False
      contains (Not x) var          = x `contains` var
      contains (BinOp _ a b) var    = a `contains` var || b `contains` var
      contains (Fun _ _) _          = False
      contains (Call _ _ f' xs) var = f' `contains` var||any (`contains` var) xs
      contains (Index a i) var      = a `contains` var || i `contains` var
      contains (Member o _) var     = o `contains` var
      contains (Arr xs) var         = any (`contains` var) xs
      contains (Obj xs) var         = any (`contains` var) (map snd xs)
      contains (AssignEx l r) var   = l `contains` var || r `contains` var
      contains (IfEx c t e) var     = any (`contains` var) [c,t,e]
      contains (Eval x) var         = x `contains` var
      contains (Thunk _ _) _        = False
  tailLoopify _ fun = do
    return fun
  
  -- | Inline a tailcalled function @f@ when:
  --
  --   * @f@ does not refer to itself; and
  --   * @f@ is defined immediately before its call site.
  --     (@let f = ... in tailcall f@)
  --
  --   Should be called *after* 'tailLoopify' but *before* trampoline for best
  --   effect.
  inlineShortJumpTailcall :: JSTrav ast => ast -> TravM ast
  inlineShortJumpTailcall ast = do
      mapJS (const True) return inl ast
    where
      inl stm@(Assign (NewVar True f) (Fun as b) tc)
        | Just (f', as') <- getTailcallInfo tc, f == f' = do
          occs <- occurrences (const True) (isEqualTo f) b
          case (occs, zipAssign (map (NewVar True) as) as' b) of
            (Never, Just b') -> return b'
            _                -> return stm
      inl stm =
        return stm
      isEqualTo v' (Exp (Var v) _) = v == v'
      isEqualTo _ _                = False
  
  -- | Extract the function being called and its argument list from a
  --   @Tailcall (Call ...)@ or @Return (Call ...)@, provided that the call is
  --   completely saturated.
  getTailcallInfo :: Stm -> Maybe (Var, [Exp])
  getTailcallInfo (Tailcall (Call 0 _ (Var f) as)) = Just (f, as)
  getTailcallInfo (Return (Call 0 _ (Var f) as))   = Just (f, as)
  getTailcallInfo _                                = Nothing
  
  -- | Assign several variables, before executing a statement.
  zipAssign :: [LHS] -> [Exp] -> Stm -> Maybe Stm
  zipAssign l r final
    | length l == length r = Just $ go l r
    | otherwise            = Nothing
    where
      go (v:vs) (x:xs) = Assign v x (go vs xs)
      go [] []         = final
      go _ _           = error "zipAssign: different number of lhs and rhs!"
  
  -- | Eliminate evaluation of vars that are guaranteed not to be thunks.
  --   Mainly useful in 'topLevelInline'.
  unevalLits :: JSTrav ast => ast -> TravM ast
  unevalLits ast = do
      lits <- foldJS (\_ _ -> True) gatherLits S.empty ast
      mapJS (const True) pure (pure . removeUpdate (`S.member` lits)) ast
    where
      gatherLits s (Stm (Assign (NewVar _ v) rhs _) _)
        | definitelyNotThunk rhs         = pure $ S.insert v s
        | Var v' <- rhs, v' `S.member` s = pure $ S.insert v s
      gatherLits s _                     = pure s
  
  -- | Inline calls to JS @eval@, @__set@, @__get@ and @__has@ and apply
  --   functions for "Haste.Foreign".
  inlineJSPrimitives :: JSTrav ast => ast -> TravM ast
  inlineJSPrimitives =
      inlineFuns >=> optimizeThunks
    where
      inlineFuns = mapJS (const True) (return . inl) return
      inl ex@(Call _ (Fast _) (Var (Foreign fn)) args) =
        case (fn, args) of
          ("eval", [Lit (LStr s)]) -> JSLit s
          ("__app0", [f])          -> Call 0 (Fast False) f []
          ("__app1", f:xs)         -> Call 0 (Fast False) f xs
          ("__app2", f:xs)         -> Call 0 (Fast False) f xs
          ("__app3", f:xs)         -> Call 0 (Fast False) f xs
          ("__app4", f:xs)         -> Call 0 (Fast False) f xs
          ("__app5", f:xs)         -> Call 0 (Fast False) f xs
          ("__get", [o, k])        -> Index o k
          ("__set", [o, k, v])     -> AssignEx (Index o k) v
          ("__has", [o, k])        -> BinOp StrictNeq (Index o k)
                                                      (JSLit "undefined")
          _                        -> ex
      inl ex =
        ex
  
  -- | Turn all assignments of the form @var v1 = e ; exp@ into @exp[v1/e]@,
  --   provided that @v1@ is not used as a known location and @e@ is either a
  --   variable or a non-string literal.
  --   Also does not inline vars into lambdas.
  --   Should go before 'tailLoopify'.
  assignToSubst :: JSTrav ast => ast -> TravM ast
  assignToSubst ast = do
      mapJS (const True) return inl ast
    where
      inl stm@(Assign (NewVar _ v) x next) | not (isKnownLoc v) = do
        case x of
          (Var _)        -> do
            -- TODO: this can be replaced by a map (to replace) and a count of
            --       remaining occurrences of v (fold)
            (c, stm') <- replaceExWithCount (isSafeForInlining) (Var v) x next
            (c', _) <- replaceExWithCount (pure True) (Var v) x next
            if c == c'
              then return stm' -- No occurrences inside lambda
              else return stm
          (Lit (LStr _)) -> return stm
          (Lit _)        -> replaceEx (pure True) (Var v) x next
          _              -> return stm
      inl stm = do
        return stm
  
  -- | Perform various trivially correct local inlinings:
  --   var x = e; return [0, e] (boxing at the end of a thunk/function)
  --    => [0, e]
  --   thunk(x) where x is non-computing and non-recursive
  --     => x
  smallStepInline :: JSTrav ast => ast -> TravM ast
  smallStepInline ast = do
      mapJS (const True) return inl ast
    where
      inl (Assign (NewVar _ v) ex (Return (Arr [l@(Lit _), Var v'])))
        | v == v' =
          return (Return (Arr [l, ex]))
      inl (Assign (NewVar _ v) ex (ThunkRet (Arr [l@(Lit _), Var v'])))
        | v == v' =
          return (ThunkRet (Arr [l, ex]))
      -- Unpack thunks which don't provide actual laziness.
      inl (Assign lhs@(NewVar _ v) t@(Thunk _ _) next)
        | Just ex <- fromThunkEx t, safeToUnThunk ex = do
          case ex of
            Thunk _ _ -> return $ Assign lhs ex next
            _         -> Assign lhs ex <$> eliminateEvalOf v next
      -- Merge @a = eval(a) ; a = eval(a)@ into a single eval.
      inl stm@(Assign (LhsExp _ (Var v1)) (Eval (Var v1')) next)
        | v1 == v1' = do
          case next of
            stm'@(Assign (LhsExp _ (Var v2)) (Eval (Var v2')) _)
              | v1 == v2 && v1' == v2' -> do
                return stm'
            _ -> do
              return stm
      inl stm =
        return stm
  
  -- | Eliminate elimination of a variable. Only use when  *absolutely certain*
  --   that the variable can never be a thunk.
  eliminateEvalOf :: JSTrav ast => Var -> ast -> TravM ast
  eliminateEvalOf v ast = mapJS (const True) return elim ast
    where
      elim (Assign (LhsExp _ (Var v')) (Eval (Var v'')) next)
        | v' == v'' && v == v' =
          return next
      elim stm =
          return stm
  
  -- | Is the given expression safe to extract from a thunk?
  --   An expression is safe to unthunk iff evaluating it will not cause
  --   computation to take place or a variable to be dereferenced.
  safeToUnThunk :: Exp -> Bool
  safeToUnThunk ex =
    case ex of
      Lit _      -> True
      JSLit l    -> isJSFunDecl l
      Fun _ _    -> True
      Thunk _ _  -> True
      Arr arr    -> all safeToUnThunk arr
      _          -> False
  
  -- | Create a mapping from (function name, argument #) to argument names
  --   for all function arguments.
  gatherArgNames :: JSTrav ast => ast -> TravM ArgMap
  gatherArgNames = foldJS (\_ _ -> True) (\m x -> pure (upd m x)) M.empty
    where
      upd m (Stm (Assign (NewVar _ v) (Fun args _) _) _) = M.insert v args m
      upd m _                                            = m
  
  flowOpts :: Stm -> TravM Stm
  flowOpts ast = do
      argvars <- gatherArgNames ast
      let info = findVarInfos argvars ast
      ast' <- mapJS (const True) (pure . removeEvals info) pure ast
      mapJS (const True) (pure . makeFastCall info) pure ast'
    where
      name (Foreign s)      = show s
      name (Internal s _ _) = show $ nameIdent s
  
      alwaysStrict infos v =
        case M.lookup v infos of
          Just info -> varStrict info == Strict
          _         -> False
  
      arityIsAlways infos v = do
        info <- M.lookup v infos
        varArity info
  
      removeEvals info (Eval ex@(Var v)) | alwaysStrict info v = ex
      removeEvals _ ex                                         = ex
  
      makeFastCall info ex@(Call _ (Normal bounce) f@(Var v) xs) =
        case arityIsAlways info v of
          Just ar'
            | length xs == ar' -> Call 0 (Fast bounce) f xs
            | length xs > ar'  -> Call 0 (Normal bounce)
                                         (Call 0 (Fast bounce) f (take ar' xs))
                                         (drop ar' xs)
          _             -> ex
      makeFastCall _ ex           = ex
  