<commit_before>module GHCVM.CodeGen.Expr where

import GHCVM.Core.CoreSyn
import GHCVM.Types.Type
import GHCVM.Types.TyCon
import GHCVM.BasicTypes.Id
import GHCVM.Prelude.PrimOp
import GHCVM.StgSyn.StgSyn
import GHCVM.BasicTypes.DataCon
import GHCVM.Utils.Panic
import GHCVM.Utils.Util (unzipWith)
import GHCVM.Util

import GHCVM.Debug
import GHCVM.CodeGen.Utils
import GHCVM.CodeGen.Monad
import GHCVM.CodeGen.Name
import GHCVM.CodeGen.Layout
import GHCVM.CodeGen.Types
import GHCVM.CodeGen.Closure
import GHCVM.CodeGen.Env
import GHCVM.CodeGen.Rts
import GHCVM.CodeGen.Constr
import GHCVM.CodeGen.Prim
import GHCVM.CodeGen.ArgRep
import GHCVM.CodeGen.LetNoEscape
import {-# SOURCE #-} GHCVM.CodeGen.Bind (cgBind)
import Codec.JVM

import Data.Monoid((<>))
import Data.Foldable(fold)
import Data.Maybe(mapMaybe)
import Control.Monad(when, forM_, unless)

cgExpr :: StgExpr -> CodeGen ()
cgExpr (StgApp fun args) = debugDoc (str "SgApp" <+> ppr fun <+> ppr args) >>
                           cgIdApp fun args
cgExpr (StgOpApp (StgPrimOp SeqOp) [StgVarArg a, _] _) = cgIdApp a []
cgExpr (StgOpApp op args ty) = debugDoc (str "StgOpApp" <+> ppr args <+> ppr ty) >>
                               cgOpApp op args ty
cgExpr (StgConApp con args) = debugDoc (str "StgConApp" <+> ppr con <+> ppr args) >>
                              cgConApp con args
-- TODO: Deal with ticks
cgExpr (StgTick t e) = cgExpr e
cgExpr (StgLit lit) = emitReturn [mkLocDirect False $ cgLit lit]
cgExpr (StgLet binds expr) = do
  cgBind binds
  cgExpr expr
cgExpr (StgLetNoEscape _ _ binds expr) =
  cgLneBinds binds expr

cgExpr (StgCase expr _ _ binder _ altType alts) =
  debugDoc (str "StgCase" <+> ppr expr <+> ppr binder <+> ppr altType) >>
  cgCase expr binder altType alts
cgExpr _ = unimplemented "cgExpr"

cgLneBinds :: StgBinding -> StgExpr -> CodeGen ()
cgLneBinds (StgNonRec binder rhs) expr = do
  (info, genBindCode) <- cgLetNoEscapeRhsBody binder rhs
  bindCode <- genBindCode
  addBinding info
  exprCode <- forkLneBody $ cgExpr expr
  let (bindLabel, argLocs) = expectJust "cgLneBinds:StgNonRec" . maybeLetNoEscape $ info
  emit $ fold (map storeDefault argLocs)
      <> letNoEscapeCodeBlocks [(bindLabel, bindCode)] exprCode

cgLneBinds (StgRec pairs) expr = do
  result <- sequence $ unzipWith cgLetNoEscapeRhsBody pairs
  let (infos, genBindCodes) = unzip result
      (labels, argLocss) = unzip $ map (expectJust "cgLneBinds:StgRec" . maybeLetNoEscape) infos
  addBindings infos
  bindCodes <- sequence genBindCodes
  exprCode <- forkLneBody $ cgExpr expr
  -- TODO: This takes the easy way out.
  -- Look here for optimization opportunities
  -- and potential source of bugs involving let-no-escape.
  -- We are banking on the fact that mututally-recursive lne bindings are
  -- rare.
  emit $ fold (fold (map (map storeDefault) argLocss))
      <> letNoEscapeCodeBlocks (zip labels bindCodes) exprCode

cgLetNoEscapeRhsBody :: Id -> StgRhs -> CodeGen (CgIdInfo, CodeGen Code)
cgLetNoEscapeRhsBody binder (StgRhsClosure _ _ _ _ _ args body)
  = cgLetNoEscapeClosure binder (nonVoidIds args) body
cgLetNoEscapeRhsBody binder (StgRhsCon _ con args)
  = cgLetNoEscapeClosure binder [] (StgConApp con args)

cgLetNoEscapeClosure
  :: Id -> [NonVoid Id] -> StgExpr -> CodeGen (CgIdInfo, CodeGen Code)
cgLetNoEscapeClosure binder args body = do
  label <- newLabel
  -- Restore the local variable count so that
  -- Those locals can be reused later on
  n <- peekNextLocal
  argLocs <- mapM newIdLoc args
  n' <- peekNextLocal
    --setNextLocal n
  let code = forkLneBody $ do
          bindArgs $ zip args argLocs
          setNextLocal n'
          cgExpr body
  return (lneIdInfo label binder argLocs, code)

cgIdApp :: Id -> [StgArg] -> CodeGen ()
cgIdApp funId [] | isVoidTy (idType funId) = emitReturn []
cgIdApp funId args = do
  dflags <- getDynFlags
  funInfo <- getCgIdInfo funId
  selfLoopInfo <- getSelfLoop<commit_msg>Fixed bug in mutually recursive lne binds<commit_after>module GHCVM.CodeGen.Expr where

import GHCVM.Core.CoreSyn
import GHCVM.Types.Type
import GHCVM.Types.TyCon
import GHCVM.BasicTypes.Id
import GHCVM.Prelude.PrimOp
import GHCVM.StgSyn.StgSyn
import GHCVM.BasicTypes.DataCon
import GHCVM.Utils.Panic
import GHCVM.Utils.Util (unzipWith)
import GHCVM.Util

import GHCVM.Debug
import GHCVM.CodeGen.Utils
import GHCVM.CodeGen.Monad
import GHCVM.CodeGen.Name
import GHCVM.CodeGen.Layout
import GHCVM.CodeGen.Types
import GHCVM.CodeGen.Closure
import GHCVM.CodeGen.Env
import GHCVM.CodeGen.Rts
import GHCVM.CodeGen.Constr
import GHCVM.CodeGen.Prim
import GHCVM.CodeGen.ArgRep
import GHCVM.CodeGen.LetNoEscape
import {-# SOURCE #-} GHCVM.CodeGen.Bind (cgBind)
import Codec.JVM

import Data.Monoid((<>))
import Data.Foldable(fold)
import Data.Maybe(mapMaybe)
import Control.Monad(when, forM_, unless)

cgExpr :: StgExpr -> CodeGen ()
cgExpr (StgApp fun args) = debugDoc (str "SgApp" <+> ppr fun <+> ppr args) >>
                           cgIdApp fun args
cgExpr (StgOpApp (StgPrimOp SeqOp) [StgVarArg a, _] _) = cgIdApp a []
cgExpr (StgOpApp op args ty) = debugDoc (str "StgOpApp" <+> ppr args <+> ppr ty) >>
                               cgOpApp op args ty
cgExpr (StgConApp con args) = debugDoc (str "StgConApp" <+> ppr con <+> ppr args) >>
                              cgConApp con args
-- TODO: Deal with ticks
cgExpr (StgTick t e) = cgExpr e
cgExpr (StgLit lit) = emitReturn [mkLocDirect False $ cgLit lit]
cgExpr (StgLet binds expr) = do
  cgBind binds
  cgExpr expr
cgExpr (StgLetNoEscape _ _ binds expr) =
  cgLneBinds binds expr

cgExpr (StgCase expr _ _ binder _ altType alts) =
  debugDoc (str "StgCase" <+> ppr expr <+> ppr binder <+> ppr altType) >>
  cgCase expr binder altType alts
cgExpr _ = unimplemented "cgExpr"

cgLneBinds :: StgBinding -> StgExpr -> CodeGen ()
cgLneBinds (StgNonRec binder rhs) expr = do
  (info, genBindCode) <- cgLetNoEscapeRhsBody binder rhs
  n' <- peekNextLocal
  bindCode <- genBindCode n'
  addBinding info
  exprCode <- forkLneBody $ cgExpr expr
  let (bindLabel, argLocs) = expectJust "cgLneBinds:StgNonRec" . maybeLetNoEscape $ info
  emit $ fold (map storeDefault argLocs)
      <> letNoEscapeCodeBlocks [(bindLabel, bindCode)] exprCode

cgLneBinds (StgRec pairs) expr = do
  result <- sequence $ unzipWith cgLetNoEscapeRhsBody pairs
  n' <- peekNextLocal
  let (infos, genBindCodes) = unzip result
      (labels, argLocss) = unzip $ map (expectJust "cgLneBinds:StgRec" . maybeLetNoEscape) infos
  addBindings infos
  bindCodes <- sequence $ ($ n') <$> genBindCodes
  exprCode <- forkLneBody $ cgExpr expr
  emit $ fold (fold (map (map storeDefault) argLocss))
      <> letNoEscapeCodeBlocks (zip labels bindCodes) exprCode

cgLetNoEscapeRhsBody :: Id -> StgRhs -> CodeGen (CgIdInfo, Int -> CodeGen Code)
cgLetNoEscapeRhsBody binder (StgRhsClosure _ _ _ _ _ args body)
  = cgLetNoEscapeClosure binder (nonVoidIds args) body
cgLetNoEscapeRhsBody binder (StgRhsCon _ con args)
  = cgLetNoEscapeClosure binder [] (StgConApp con args)

cgLetNoEscapeClosure
  :: Id -> [NonVoid Id] -> StgExpr -> CodeGen (CgIdInfo, Int -> CodeGen Code)
cgLetNoEscapeClosure binder args body = do
  label <- newLabel
  n <- peekNextLocal
  argLocs <- mapM newIdLoc args
  let code n' = forkLneBody $ do
        bindArgs $ zip args argLocs
        setNextLocal n'
        cgExpr body
  return (lneIdInfo label binder argLocs, code)

cgIdApp :: Id -> [StgArg] -> CodeGen ()
cgIdApp funId [] | isVoidTy (idType funId) = emitReturn []
cgIdApp funId args = do
  dflags <- getDynFlags
  funInfo <- getCgIdInfo funId
  selfLoopInfo <- getSelfLoop
{'fraction_non_alphanumeric': 0.09307844429795649, 'fraction_numerical': 0.0, 'mean_word_length': 3.7800882167611847, 'pattern_counts': {'":': 0, '<': 55, '<?xml version=': 0, '>': 66, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7195871', 'n_tokens_mistral': 2835, 'n_tokens_neox': 2731, 'n_words': 761}