content: -- | This module performs basic inlining of known functions
  module Language.PureScript.CoreImp.Optimizer.Inliner
    ( inlineVariables
    , inlineCommonValues
    , inlineCommonOperators
    , inlineFnComposition
    , inlineFnIdentity
    , inlineUnsafeCoerce
    , inlineUnsafePartial
    , etaConvert
    , unThunk
    , evaluateIifes
    ) where
  
  import Prelude
  
  import Control.Monad.Supply.Class (MonadSupply, freshName)
  
  import Data.Either (rights)
  import Data.Maybe (fromMaybe)
  import Data.Text (Text)
  import Data.Text qualified as T
  
  import Language.PureScript.Names (ModuleName)
  import Language.PureScript.PSString (PSString, mkString)
  import Language.PureScript.CoreImp.AST (AST(..), BinaryOperator(..), InitializerEffects(..), UnaryOperator(..), everywhere, everywhereTopDown, everywhereTopDownM, getSourceSpan)
  import Language.PureScript.CoreImp.Optimizer.Common (pattern Ref, applyAll, isReassigned, isRebound, isUpdated, removeFromBlock, replaceIdent, replaceIdents)
  import Language.PureScript.AST (SourceSpan(..))
  import Language.PureScript.Constants.Libs qualified as C
  import Language.PureScript.Constants.Prim qualified as C
  
  -- TODO: Potential bug:
  -- Shouldn't just inline this case: { var x = 0; x.toFixed(10); }
  -- Needs to be: { 0..toFixed(10); }
  -- Probably needs to be fixed in pretty-printer instead.
  shouldInline :: AST -> Bool
  shouldInline (Var _ _) = True
  shouldInline (ModuleAccessor _ _ _) = True
  shouldInline (NumericLiteral _ _) = True
  shouldInline (StringLiteral _ _) = True
  shouldInline (BooleanLiteral _ _) = True
  shouldInline (Indexer _ index val) = shouldInline index && shouldInline val
  shouldInline _ = False
  
  etaConvert :: AST -> AST
  etaConvert = everywhere convert
    where
    convert :: AST -> AST
    convert (Block ss [Return _ (App _ (Function _ Nothing idents block@(Block _ body)) args)])
      | all shouldInline args &&
        not (any ((`isRebound` block) . Var Nothing) idents) &&
        not (any (`isRebound` block) args)
        = Block ss (map (replaceIdents (zip idents args)) body)
    convert (Function _ Nothing [] (Block _ [Return _ (App _ fn [])])) = fn
    convert js = js
  
  unThunk :: AST -> AST
  unThunk = everywhere convert
    where
    convert :: AST -> AST
    convert (Block ss []) = Block ss []
    convert (Block ss jss) =
      case last jss of
        Return _ (App _ (Function _ Nothing [] (Block _ body)) []) -> Block ss $ init jss ++ body
        _ -> Block ss jss
    convert js = js
  
  evaluateIifes :: AST -> AST
  evaluateIifes = everywhere convert
    where
    convert :: AST -> AST
    convert (App _ (Function _ Nothing [] (Block _ [Return _ ret])) []) = ret
    convert (App _ (Function _ Nothing idents (Block _ [Return ss ret])) [])
      | not (any (`isReassigned` ret) idents) = replaceIdents (map (, Var ss C.S_undefined) idents) ret
    convert js = js
  
  inlineVariables :: AST -> AST
  inlineVariables = everywhere $ removeFromBlock go
    where
    go :: [AST] -> [AST]
    go [] = []
    go (VariableIntroduction _ var (Just (_, js)) : sts)
      | shouldInline js && not (any (isReassigned var) sts) && not (any (isRebound js) sts) && not (any (isUpdated var) sts) =
        go (map (replaceIdent var js) sts)
    go (s:sts) = s : go sts
  
  inlineCommonValues :: (AST -> AST) -> AST -> AST
  inlineCommonValues expander = everywhere convert
    where
    convert :: AST -> AST
    convert (expander -> App ss (Ref fn) [Ref dict])
      | dict `elem` [C.P_semiringNumber, C.P_semiringInt], C.P_zero <- fn = NumericLiteral ss (Left 0)
      | dict `elem` [C.P_semiringNumber, C.P_semiringInt], C.P_one <- fn = NumericLiteral ss (Left 1)
      | C.P_boundedBoolean <- dict, C.P_bottom <- fn = BooleanLiteral ss False
      | C.P_boundedBoolean <- dict, C.P_top <- fn = BooleanLiteral ss True
    convert (App ss (expander -> App _ (Ref C.P_negate) [Ref C.P_ringInt]) [x])
      = Binary ss BitwiseOr (Unary ss Negate x) (NumericLiteral ss (Left 0))
    convert (App ss (App _ (expander -> App _ (Ref fn) [Ref dict]) [x]) [y])
      | C.P_semiringInt <- dict, C.P_add <- fn = intOp ss Add x y
      | C.P_semiringInt <- dict, C.P_mul <- fn = intOp ss Multiply x y
      | C.P_ringInt <- dict, C.P_sub <- fn = intOp ss Subtract x y
    convert other = other
    intOp ss op x y = Binary ss BitwiseOr (Binary ss op x y) (NumericLiteral ss (Left 0))
  
  inlineCommonOperators :: (AST -> AST) -> AST -> AST
  inlineCommonOperators expander = everywhereTopDown $ applyAll $
    [ binary C.P_semiringNumber C.P_add Add
    , binary C.P_semiringNumber C.P_mul Multiply
  
    , binary C.P_ringNumber C.P_sub Subtract
    , unary  C.P_ringNumber C.P_negate Negate
  
    , binary C.P_euclideanRingNumber C.P_div Divide
  
    , binary C.P_eqNumber C.P_eq EqualTo
    , binary C.P_eqNumber C.P_notEq NotEqualTo
    , binary C.P_eqInt C.P_eq EqualTo
    , binary C.P_eqInt C.P_notEq NotEqualTo
    , binary C.P_eqString C.P_eq EqualTo
    , binary C.P_eqString C.P_notEq NotEqualTo
    , binary C.P_eqChar C.P_eq EqualTo
    , binary C.P_eqChar C.P_notEq NotEqualTo
    , binary C.P_eqBoolean C.P_eq EqualTo
    , binary C.P_eqBoolean C.P_notEq NotEqualTo
  
    , binary C.P_ordBoolean C.P_lessThan LessThan
    , binary C.P_ordBoolean C.P_lessThanOrEq LessThanOrEqualTo
    , binary C.P_ordBoolean C.P_greaterThan GreaterThan
    , binary C.P_ordBoolean C.P_greaterThanOrEq GreaterThanOrEqualTo
    , binary C.P_ordChar C.P_lessThan LessThan
    , binary C.P_ordChar C.P_lessThanOrEq LessThanOrEqualTo
    , binary C.P_ordChar C.P_greaterThan GreaterThan
    , binary C.P_ordChar C.P_greaterThanOrEq GreaterThanOrEqualTo
    , binary C.P_ordInt C.P_lessThan LessThan
    , binary C.P_ordInt C.P_lessThanOrEq LessThanOrEqualTo
    , binary C.P_ordInt C.P_greaterThan GreaterThan
    , binary C.P_ordInt C.P_greaterThanOrEq GreaterThanOrEqualTo
    , binary C.P_ordNumber C.P_lessThan LessThan
    , binary C.P_ordNumber C.P_lessThanOrEq LessThanOrEqualTo
    , binary C.P_ordNumber C.P_greaterThan GreaterThan
    , binary C.P_ordNumber C.P_greaterThanOrEq GreaterThanOrEqualTo
    , binary C.P_ordString C.P_lessThan LessThan
    , binary C.P_ordString C.P_lessThanOrEq LessThanOrEqualTo
    , binary C.P_ordString C.P_greaterThan GreaterThan
    , binary C.P_ordString C.P_greaterThanOrEq GreaterThanOrEqualTo
  
    , binary C.P_semigroupString C.P_append Add
  
    , binary C.P_heytingAlgebraBoolean C.P_conj And
    , binary C.P_heytingAlgebraBoolean C.P_disj Or
    , unary  C.P_heytingAlgebraBoolean C.P_not Not
  
    , binary' C.P_or BitwiseOr
    , binary' C.P_and BitwiseAnd
    , binary' C.P_xor BitwiseXor
    , binary' C.P_shl ShiftLeft
    , binary' C.P_shr ShiftRight
    , binary' C.P_zshr ZeroFillShiftRight
    , unary'  C.P_complement BitwiseNot
  
    , inlineNonClassFunction (isModFnWithDict C.P_unsafeIndex) $ flip (Indexer Nothing)
    ] ++
    [ fn | i <- [0..10], fn <- [ mkFn i, runFn i ] ] ++
    [ fn | i <- [0..10], fn <- [ mkEffFn C.P_mkEffFn i, runEffFn C.P_runEffFn i ] ] ++
    [ fn | i <- [0..10], fn <- [ mkEffFn C.P_mkEffectFn i, runEffFn C.P_runEffectFn i ] ] ++
    [ fn | i <- [0..10], fn <- [ mkEffFn C.P_mkSTFn i, runEffFn C.P_runSTFn i ] ]
    where
    binary :: (ModuleName, PSString) -> (ModuleName, PSString) -> BinaryOperator -> AST -> AST
    binary dict fn op = convert where
      convert :: AST -> AST
      convert (App ss (App _ (expander -> App _ (Ref fn') [Ref dict']) [x]) [y]) | dict == dict', fn == fn' = Binary ss op x y
      convert other = other
    binary' :: (ModuleName, PSString) -> BinaryOperator -> AST -> AST
    binary' fn op = convert where
      convert :: AST -> AST
      convert (App ss (App _ (Ref fn') [x]) [y]) | fn == fn' = Binary ss op x y
      convert other = other
    unary :: (ModuleName, PSString) -> (ModuleName, PSString) -> UnaryOperator -> AST -> AST
    unary dict fn op = convert where
      convert :: AST -> AST
      convert (App ss (expander -> App _ (Ref fn') [Ref dict']) [x]) | dict == dict', fn == fn' = Unary ss op x
      convert other = other
    unary' :: (ModuleName, PSString) -> UnaryOperator -> AST -> AST
    unary' fn op = convert where
      convert :: AST -> AST
      convert (App ss (Ref fn') [x]) | fn == fn' = Unary ss op x
      convert other = other
  
    mkFn :: Int -> AST -> AST
    mkFn = mkFn' C.P_mkFn $ \ss1 ss2 ss3 args js ->
      Function ss1 Nothing args (Block ss2 [Return ss3 js])
  
    mkEffFn :: (ModuleName, PSString) -> Int -> AST -> AST
    mkEffFn mkFn_ = mkFn' mkFn_ $ \ss1 ss2 ss3 args js ->
      Function ss1 Nothing args (Block ss2 [Return ss3 (App ss3 js [])])
  
    mkFn' :: (ModuleName, PSString) -> (Maybe SourceSpan -> Maybe SourceSpan -> Maybe SourceSpan -> [Text] -> AST -> AST) -> Int -> AST -> AST
    mkFn' mkFn_ res 0 = convert where
      convert :: AST -> AST
      convert (App _ (Ref mkFnN) [Function s1 Nothing [_] (Block s2 [Return s3 js])]) | isNFn mkFn_ 0 mkFnN =
        res s1 s2 s3 [] js
      convert other = other
    mkFn' mkFn_ res n = convert where
      convert :: AST -> AST
      convert orig@(App ss (Ref mkFnN) [fn]) | isNFn mkFn_ n mkFnN =
        case collectArgs n [] fn of
          Just (args, [Return ss' ret]) -> res ss ss ss' args ret
          _ -> orig
      convert other = other
      collectArgs :: Int -> [Text] -> AST -> Maybe ([Text], [AST])
      collectArgs 1 acc (Function _ Nothing [oneArg] (Block _ js)) | length acc == n - 1 = Just (reverse (oneArg : acc), js)
      collectArgs m acc (Function _ Nothing [oneArg] (Block _ [Return _ ret])) = collectArgs (m - 1) (oneArg : acc) ret
      collectArgs _ _   _ = Nothing
  
    isNFn :: (ModuleName, PSString) -> Int -> (ModuleName, PSString) -> Bool
    isNFn prefix n fn = fmap (<> mkString (T.pack $ show n)) prefix == fn
  
    runFn :: Int -> AST -> AST
    runFn = runFn' C.P_runFn App
  
    runEffFn :: (ModuleName, PSString) -> Int -> AST -> AST
    runEffFn runFn_ = runFn' runFn_ $ \ss fn acc ->
      Function ss Nothing [] (Block ss [Return ss (App ss fn acc)])
  
    runFn' :: (ModuleName, PSString) -> (Maybe SourceSpan -> AST -> [AST] -> AST) -> Int -> AST -> AST
    runFn' runFn_ res n = convert where
      convert :: AST -> AST
      convert js = fromMaybe js $ go n [] js
  
      go :: Int -> [AST] -> AST -> Maybe AST
      go 0 acc (App ss (Ref runFnN) [fn]) | isNFn runFn_ n runFnN && length acc == n =
        Just $ res ss fn acc
      go m acc (App _ lhs [arg]) = go (m - 1) (arg : acc) lhs
      go _ _   _ = Nothing
  
    inlineNonClassFunction :: (AST -> Bool) -> (AST -> AST -> AST) -> AST -> AST
    inlineNonClassFunction p f = convert where
      convert :: AST -> AST
      convert (App _ (App _ op' [x]) [y]) | p op' = f x y
      convert other = other
  
    isModFnWithDict :: (ModuleName, PSString) -> AST -> Bool
    isModFnWithDict fn (App _ (Ref fn') [Var _ _]) = fn == fn'
    isModFnWithDict _ _ = False
  
  -- (f <<< g $ x) = f (g x)
  -- (f <<< g)     = \x -> f (g x)
  inlineFnComposition :: forall m. MonadSupply m => (AST -> AST) -> AST -> m AST
  inlineFnComposition expander = everywhereTopDownM convert
    where
    convert :: AST -> m AST
    convert (App s1 (App s2 (App _ (expander -> App _ (Ref fn) [Ref C.P_semigroupoidFn]) [x]) [y]) [z])
      | C.P_compose <- fn = return $ App s1 x [App s2 y [z]]
      | C.P_composeFlipped <- fn = return $ App s2 y [App s1 x [z]]
    convert app@(App ss (App _ (expander -> App _ (Ref fn) [Ref C.P_semigroupoidFn]) _) _)
      | fn `elem` [C.P_compose, C.P_composeFlipped] = mkApps ss <$> goApps app <*> freshName
    convert other = return other
  
    mkApps :: Maybe SourceSpan -> [Either AST (Text, AST)] -> Text -> AST
    mkApps ss fns a = App ss (Function ss Nothing [] (Block ss $ vars <> [Return Nothing comp])) []
      where
      vars = uncurry (VariableIntroduction ss) . fmap (Just . (UnknownEffects, )) <$> rights fns
      comp = Function ss Nothing [a] (Block ss [Return Nothing apps])
      apps = foldr (\fn acc -> App ss (mkApp fn) [acc]) (Var ss a) fns
  
    mkApp :: Either AST (Text, AST) -> AST
    mkApp = either id $ \(name, arg) -> Var (getSourceSpan arg) name
  
    goApps :: AST -> m [Either AST (Text, AST)]
    goApps (App _ (App _ (expander -> App _ (Ref fn) [Ref C.P_semigroupoidFn]) [x]) [y])
      | C.P_compose <- fn = mappend <$> goApps x <*> goApps y
      | C.P_composeFlipped <- fn = mappend <$> goApps y <*> goApps x
    goApps app@App {} = pure . Right . (,app) <$> freshName
    goApps other = pure [Left other]
  
  inlineFnIdentity :: (AST -> AST) -> AST -> AST
  inlineFnIdentity expander = everywhereTopDown convert
    where
    convert :: AST -> AST
    convert (App _ (expander -> App _ (Ref C.P_identity) [Ref C.P_categoryFn]) [x]) = x
    convert other = other
  
  inlineUnsafeCoerce :: AST -> AST
  inlineUnsafeCoerce = everywhereTopDown convert where
    convert (App _ (Ref C.P_unsafeCoerce) [ comp ]) = comp
    convert other = other
  
  inlineUnsafePartial :: AST -> AST
  inlineUnsafePartial = everywhereTopDown convert where
    convert (App ss (Ref C.P_unsafePartial) [ comp ])
      -- Apply to undefined here, the application should be optimized away
      -- if it is safe to do so
      = App ss comp [ Var ss C.S_undefined ]
    convert other = other
  