Fix where inlineSat is called.

IntersectMBO · Mar 21, 2023 · c49f792 · c49f792
1 parent ab03640
commit c49f792
Show file tree

Hide file tree

Showing 14 changed files with 144 additions and 515 deletions.
diff --git a/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/CallSiteInline.hs b/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/CallSiteInline.hs
@@ -155,27 +155,6 @@ computeArity = \case
     -- Whenever we encounter a body that is not a lambda or type abstraction, we are done counting
     tm                -> ([],tm)
 
--- | Inline fully applied functions iff the body of the function is `acceptable`.
-considerInline :: forall tyname name uni fun ann. InliningConstraints tyname name uni fun
-    => Term tyname name uni fun ann -- the variable that is a function
-    -> InlineM tyname name uni fun ann (Term tyname name uni fun ann)
-considerInline v@(Var _ann n) = do
-  -- look up the variable in the `CalledVar` map
-  varInfo <- gets (lookupCalled n)
-  case varInfo of
-      -- if it's not in the map, it's not a function, don't inline.
-      Nothing -> pure v
-      Just info -> do
-        isAcceptable <- acceptable (calledVarBody info)
-      -- if the size and cost are not acceptable, don't inline
-        if not isAcceptable then pure v
-        -- if the size and cost are acceptable, then check if it's fully applied
-        -- See note [Identifying fully applied call sites].
-        else
-          inlineSat v
-considerInline _notVar = -- this should not happen
-  Prelude.error  "considerInline: should be a variable."
-
 -- | A term or type argument.
 data Args tyname name uni fun ann =
   MkTermArg (Term tyname name uni fun ann)
@@ -219,56 +198,38 @@ enoughArgs lamOrder argsOrder =
     (MkType, MkTypeArg _) -> enoughArgs (init lamOrder) (init argsOrder)
     _                     -> False
 
--- | Inline fully applied functions. See note [Identifying fully applied call sites].
+-- | Inline fully applied functions iff the body of the function is `acceptable`.
 inlineSat :: forall tyname name uni fun ann. InliningConstraints tyname name uni fun
     => Term tyname name uni fun ann -- ^ The `body` of the `Let` term.
     -> InlineM tyname name uni fun ann (Term tyname name uni fun ann)
--- If the term is a term application, see if we it's applying to something that we may inline
-inlineSat appTerm@(Apply _varAnn _fun _arg) = do
-    -- collect all the arguments of the term being applied to
-    let argsAppliedTo = fst $ collectArgs appTerm
-        args = snd $ collectArgs appTerm
-    case argsAppliedTo of
-      -- if it is a `Var` that is being applied to, check to see if it's fully applied
-      Var _ann name -> do
-        maybeVarInfo <- gets (lookupCalled name)
-        case maybeVarInfo of
-          -- the variable is not in the map that contains all the in-scope functions, this shouldn't
-          -- happen? TODO maybe error out instead?
-          -- Have checked that with the current tests it doesn't happen
-          Nothing -> forMOf termSubterms appTerm inlineSat
-          Just varInfo -> do
-            if enoughArgs (arity varInfo) (map fst args) then do
-              -- if the `Var` is fully applied (over-application is allowed) then inline it
-              let inlinedTm = mkApps (calledVarDef varInfo) args
-              forMOf termSubterms inlinedTm inlineSat
-            -- otherwise just keep going
-            else forMOf termSubterms appTerm inlineSat
-          -- if the term being applied is not a `Var`, don't inline, but keep checking
-      v -> forMOf termSubterms v inlineSat -- keep checking all subterms
-inlineSat tyInstTerm@(TyInst _varAnn _fun _arg) = do
-  -- collect all the arguments of the term being applied to
-    let argsAppliedTo = fst $ collectArgs tyInstTerm
-        args = snd $ collectArgs tyInstTerm
-    case argsAppliedTo of
-      -- if it is a `Var` that is being applied to, check to see if it's fully applied
-      Var _ann name -> do
-        maybeVarInfo <- gets (lookupCalled name)
-        case maybeVarInfo of
-          Nothing -> forMOf termSubterms tyInstTerm inlineSat
-          Just varInfo -> do
-            if enoughArgs (arity varInfo) (map fst args) then do
-              -- if the `Var` is fully applied (over-application is allowed) then inline it
-              let inlinedTm = mkApps (calledVarDef varInfo) args
-              forMOf termSubterms inlinedTm inlineSat
-            -- otherwise just keep going
-            else forMOf termSubterms tyInstTerm inlineSat
-          -- if the term being applied is not a `Var`, don't inline but keep checking the subterms
-      v -> forMOf termSubterms v inlineSat
-inlineSat letTm@(Let _ _ _bds _letBody) =
-    -- recursive or not, the bindings of this let term *may* contain a saturated function,
-    -- so we need to check all the bindings and also the body
-    -- `PlutusIR.Core.Plated.termSubterms` gives all that
-    forMOf termSubterms letTm inlineSat
-inlineSat tm =
-  forMOf termSubterms tm inlineSat
+inlineSat appOrTyInstTm =
+    case appOrTyInstTm of
+      -- If the term is a term or type application, check it's applying to a var that we may inline
+      Apply _varAnn _fun _arg  -> go appOrTyInstTm
+      TyInst _varAnn _fun _arg -> go appOrTyInstTm
+      -- otherwise, check all subterms
+      _                        -> forMOf termSubterms appOrTyInstTm inlineSat
+      where
+        go tm = do
+            -- collect all the arguments of the term being applied to
+            let argsAppliedTo = fst $ collectArgs tm
+                args = snd $ collectArgs tm
+            case argsAppliedTo of
+              -- if it is a `Var` that is being applied to, check to see if it's fully applied
+              Var _ann name -> do
+                maybeVarInfo <- gets (lookupCalled name)
+                case maybeVarInfo of
+                  -- the variable is not in the map that contains all the in-scope functions,
+                  -- this shouldn't happen? TODO maybe error out instead?
+                  -- Have checked that with the current tests it doesn't happen
+                  Nothing -> forMOf termSubterms tm inlineSat
+                  Just varInfo -> do
+                    isAcceptable <- acceptable (calledVarBody varInfo)
+                    if isAcceptable && enoughArgs (arity varInfo) (map fst args) then do
+                      -- if the `Var` is fully applied (over-application is allowed) then inline it
+                      let inlinedTm = mkApps (calledVarDef varInfo) args
+                      forMOf termSubterms inlinedTm inlineSat
+                    -- otherwise just keep going
+                    else forMOf termSubterms tm inlineSat
+                  -- if the term being applied is not a `Var`, don't inline, but keep checking
+              v -> forMOf termSubterms v inlineSat -- keep checking all subterms
diff --git a/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/UnconditionalInline.hs b/plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/UnconditionalInline.hs
@@ -162,14 +162,15 @@ processTerm = handleTerm <=< traverseOf termSubtypes applyTypeSubstitution where
         Term tyname name uni fun ann
         -> InlineM tyname name uni fun ann (Term tyname name uni fun ann)
     handleTerm = \case
-        v@(Var _ n) -> do
-            inSubMap <- substName n
-            case inSubMap of
-                -- If it's not in the substitution map, check if it's saturated
-                Nothing -> do
-                    considerInline v
-                -- If it's in the substitution map, do the substitution
-                Just var -> pure var
+        v@(Var _ n) -> fromMaybe v <$> substName n
+        -- we need to check at an `Apply` note if there is a saturated function to be inlined.
+        appTerm@(Apply _varAnn _fun _arg) -> do
+            inlinedSat <- inlineSat appTerm
+            forMOf termSubterms inlinedSat processTerm
+        -- we need to check at a `TyInst` note if there is a saturated function to be inlined.
+        tyInstTerm@(TyInst _varAnn _fun _arg) -> do
+            inlinedSat <- inlineSat tyInstTerm
+            forMOf termSubterms inlinedSat processTerm
         Let ann NonRec bs t -> do
             -- Process bindings, eliminating those which will be inlined unconditionally,
             -- and accumulating the new substitutions