Optimize single-argument unknown function invocations with known non-…

…function target types (#3270) * Closes #2935 * Implements the `ConvertUnaryCalls` optimization which replaces generic calls to unknown unary functions with known non-function target types by direct closure calls * Adds an optimization phase in JuvixTree * Adds options to the JuvixTree pipeline * Fixes a bug in Rust code generation for direct closure calls
anoma · Jan 13, 2025 · 6055b0d · 6055b0d
1 parent df0c017
commit 6055b0d
Show file tree

Hide file tree

Showing 27 changed files with 233 additions and 29 deletions.
diff --git a/app/Commands/Dev/Tree/Compile/Asm.hs b/app/Commands/Dev/Tree/Compile/Asm.hs
@@ -14,9 +14,11 @@ runCommand opts = do
   outFile <- getOutputFile FileExtJuvixAsm (Just inputFile) moutputFile
   mainFile <- getMainFile (Just inputFile)
   tab :: InfoTable <- readFile mainFile >>= getRight . Tree.runParser mainFile
+  ep <- getEntryPoint (Just inputFile)
   res <-
     getRight
       . run
+      . runReader ep
       . runError @JuvixError
       $ treeToAsm tab
   writeFileEnsureLn outFile (Asm.ppPrint res res)
diff --git a/app/Commands/Dev/Tree/Read.hs b/app/Commands/Dev/Tree/Read.hs
@@ -3,6 +3,7 @@ module Commands.Dev.Tree.Read where
 import Commands.Base
 import Commands.Dev.Tree.Read.Options
 import Juvix.Compiler.Tree.Data.InfoTable qualified as Tree
+import Juvix.Compiler.Tree.Options qualified as TreeOptions
 import Juvix.Compiler.Tree.Pretty qualified as Tree
 import Juvix.Compiler.Tree.Transformation qualified as Tree
 import Juvix.Compiler.Tree.Translation.FromSource qualified as Tree
@@ -15,7 +16,7 @@ runCommand opts = do
   case Tree.runParser afile s of
     Left err -> exitJuvixError (JuvixError err)
     Right tab -> do
-      r <- runError @JuvixError (Tree.applyTransformations (project opts ^. treeReadTransformations) tab)
+      r <- runReader TreeOptions.defaultOptions $ runError @JuvixError (Tree.applyTransformations (project opts ^. treeReadTransformations) tab)
       case r of
         Left err -> exitJuvixError (JuvixError err)
         Right tab' -> do

diff --git a/src/Juvix/Compiler/Backend/Rust/Translation/FromReg.hs b/src/Juvix/Compiler/Backend/Rust/Translation/FromReg.hs
@@ -331,7 +331,7 @@ fromRegInstr info = \case
           ]
         Reg.CallClosure vr ->
           [ stmtAssign
-              "funid"
+              "(funid, args)"
               ( mkCall
                   "mem.call_closure"
                   [fromVarRef vr, mkArray ((map fromValue _instrTailCallArgs))]

diff --git a/src/Juvix/Compiler/Pipeline.hs b/src/Juvix/Compiler/Pipeline.hs
@@ -300,10 +300,10 @@ coreToVampIR = Core.toStored >=> storedCoreToVampIR
 -- Other workflows
 --------------------------------------------------------------------------------
 
-treeToAsm :: (Member (Error JuvixError) r) => Tree.InfoTable -> Sem r Asm.InfoTable
+treeToAsm :: (Members '[Error JuvixError, Reader EntryPoint] r) => Tree.InfoTable -> Sem r Asm.InfoTable
 treeToAsm = Tree.toAsm >=> return . Asm.fromTree
 
-treeToCairoAsm :: (Member (Error JuvixError) r) => Tree.InfoTable -> Sem r Asm.InfoTable
+treeToCairoAsm :: (Members '[Error JuvixError, Reader EntryPoint] r) => Tree.InfoTable -> Sem r Asm.InfoTable
 treeToCairoAsm = Tree.toCairoAsm >=> return . Asm.fromTree
 
 treeToReg :: (Members '[Error JuvixError, Reader EntryPoint] r) => Tree.InfoTable -> Sem r Reg.InfoTable

diff --git a/src/Juvix/Compiler/Tree/Data/TransformationId.hs b/src/Juvix/Compiler/Tree/Data/TransformationId.hs
@@ -8,6 +8,8 @@ data TransformationId
   = IdentityTrans
   | IdentityU
   | IdentityD
+  | ConvertUnaryCalls
+  | OptPhaseMain
   | Apply
   | FilterUnreachable
   | Validate
@@ -24,20 +26,22 @@ data PipelineId
 type TransformationLikeId = TransformationLikeId' TransformationId PipelineId
 
 toNockmaTransformations :: [TransformationId]
-toNockmaTransformations = [Validate, Apply, FilterUnreachable]
+toNockmaTransformations = [Validate, OptPhaseMain, Apply, FilterUnreachable]
 
 toAsmTransformations :: [TransformationId]
-toAsmTransformations = [Validate, CheckNoAnoma, CheckNoByteArray]
+toAsmTransformations = [Validate, CheckNoAnoma, CheckNoByteArray, OptPhaseMain]
 
 toCairoAsmTransformations :: [TransformationId]
-toCairoAsmTransformations = [Validate, Apply, FilterUnreachable]
+toCairoAsmTransformations = [Validate, OptPhaseMain, Apply, FilterUnreachable]
 
 instance TransformationId' TransformationId where
   transformationText :: TransformationId -> Text
   transformationText = \case
     IdentityTrans -> strIdentity
     IdentityU -> strIdentityU
     IdentityD -> strIdentityD
+    ConvertUnaryCalls -> strConvertUnaryCalls
+    OptPhaseMain -> strOptPhaseMain
     Apply -> strApply
     FilterUnreachable -> strFilterUnreachable
     Validate -> strValidate

diff --git a/src/Juvix/Compiler/Tree/Data/TransformationId/Strings.hs b/src/Juvix/Compiler/Tree/Data/TransformationId/Strings.hs
@@ -20,6 +20,12 @@ strIdentityU = "identity-umap"
 strIdentityD :: Text
 strIdentityD = "identity-dmap"
 
+strConvertUnaryCalls :: Text
+strConvertUnaryCalls = "convert-unary-calls"
+
+strOptPhaseMain :: Text
+strOptPhaseMain = "opt-phase-main"
+
 strApply :: Text
 strApply = "apply"
 

diff --git a/src/Juvix/Compiler/Tree/Extra/Base.hs b/src/Juvix/Compiler/Tree/Extra/Base.hs
@@ -298,7 +298,9 @@ destruct = \case
 
       mkBranchChild :: CaseBranch -> NodeChild
       mkBranchChild CaseBranch {..} =
-        (if _caseBranchSave then oneTempVar (TempVar Nothing Nothing) else noTempVar) _caseBranchBody
+        (if _caseBranchSave then oneTempVar (TempVar Nothing Nothing indTy) else noTempVar) _caseBranchBody
+        where
+          indTy = TyInductive (TypeInductive _nodeCaseInductive)
 
       mkBranches :: [CaseBranch] -> [Node] -> [CaseBranch]
       mkBranches = zipWithExact (flip (set caseBranchBody))

diff --git a/src/Juvix/Compiler/Tree/Extra/Type.hs b/src/Juvix/Compiler/Tree/Extra/Type.hs
@@ -16,6 +16,21 @@ import Juvix.Compiler.Tree.Extra.Type.Base
 import Juvix.Compiler.Tree.Language.Base
 import Juvix.Compiler.Tree.Pretty
 
+isConcreteAtomType :: Type -> Bool
+isConcreteAtomType = \case
+  TyConstr {} -> True
+  TyInductive {} -> True
+  TyInteger {} -> True
+  TyBool {} -> True
+  TyString -> True
+  TyField -> True
+  TyByteArray -> True
+  TyRandomGenerator -> True
+  TyUnit -> True
+  TyVoid -> True
+  TyDynamic -> False
+  TyFun {} -> False
+
 unfoldType :: Type -> ([Type], Type)
 unfoldType ty = (typeArgs ty, typeTarget ty)
 

diff --git a/src/Juvix/Compiler/Tree/Language.hs b/src/Juvix/Compiler/Tree/Language.hs
@@ -1,12 +1,14 @@
 module Juvix.Compiler.Tree.Language
   ( module Juvix.Compiler.Tree.Language,
     module Juvix.Compiler.Tree.Language.Base,
+    module Juvix.Compiler.Tree.Language.Type,
     module Juvix.Compiler.Tree.Language.Builtins,
   )
 where
 
 import Juvix.Compiler.Tree.Language.Base
 import Juvix.Compiler.Tree.Language.Builtins
+import Juvix.Compiler.Tree.Language.Type
 
 -- | Function call type
 data CallType
@@ -177,7 +179,8 @@ data CaseBranch = CaseBranch
 
 data TempVar = TempVar
   { _tempVarName :: Maybe Text,
-    _tempVarLocation :: Maybe Location
+    _tempVarLocation :: Maybe Location,
+    _tempVarType :: Type
   }
 
 data NodeSave = NodeSave

diff --git a/src/Juvix/Compiler/Tree/Pipeline.hs b/src/Juvix/Compiler/Tree/Pipeline.hs
@@ -4,14 +4,15 @@ module Juvix.Compiler.Tree.Pipeline
   )
 where
 
+import Juvix.Compiler.Pipeline.EntryPoint
 import Juvix.Compiler.Tree.Data.InfoTable
 import Juvix.Compiler.Tree.Transformation
 
-toNockma :: (Member (Error JuvixError) r) => InfoTable -> Sem r InfoTable
-toNockma = applyTransformations toNockmaTransformations
+toNockma :: (Members '[Error JuvixError, Reader EntryPoint] r) => InfoTable -> Sem r InfoTable
+toNockma = mapReader fromEntryPoint . applyTransformations toNockmaTransformations
 
-toAsm :: (Member (Error JuvixError) r) => InfoTable -> Sem r InfoTable
-toAsm = applyTransformations toAsmTransformations
+toAsm :: (Members '[Error JuvixError, Reader EntryPoint] r) => InfoTable -> Sem r InfoTable
+toAsm = mapReader fromEntryPoint . applyTransformations toAsmTransformations
 
-toCairoAsm :: (Member (Error JuvixError) r) => InfoTable -> Sem r InfoTable
-toCairoAsm = applyTransformations toCairoAsmTransformations
+toCairoAsm :: (Members '[Error JuvixError, Reader EntryPoint] r) => InfoTable -> Sem r InfoTable
+toCairoAsm = mapReader fromEntryPoint . applyTransformations toCairoAsmTransformations
diff --git a/src/Juvix/Compiler/Tree/Transformation.hs b/src/Juvix/Compiler/Tree/Transformation.hs
@@ -14,16 +14,20 @@ import Juvix.Compiler.Tree.Transformation.CheckNoAnoma
 import Juvix.Compiler.Tree.Transformation.CheckNoByteArray
 import Juvix.Compiler.Tree.Transformation.FilterUnreachable
 import Juvix.Compiler.Tree.Transformation.IdentityTrans
+import Juvix.Compiler.Tree.Transformation.Optimize.ConvertUnaryCalls
+import Juvix.Compiler.Tree.Transformation.Optimize.Phase.Main
 import Juvix.Compiler.Tree.Transformation.Validate
 
-applyTransformations :: forall r. (Member (Error JuvixError) r) => [TransformationId] -> InfoTable -> Sem r InfoTable
+applyTransformations :: forall r. (Members '[Error JuvixError, Reader Options] r) => [TransformationId] -> InfoTable -> Sem r InfoTable
 applyTransformations ts tbl = foldM (flip appTrans) tbl ts
   where
     appTrans :: TransformationId -> InfoTable -> Sem r InfoTable
     appTrans = \case
       IdentityTrans -> return . identity
       IdentityU -> return . identityU
       IdentityD -> return . identityD
+      ConvertUnaryCalls -> return . convertUnaryCalls
+      OptPhaseMain -> optimize
       Apply -> return . computeApply
       FilterUnreachable -> return . filterUnreachable
       Validate -> mapError (JuvixError @TreeError) . validate

diff --git a/src/Juvix/Compiler/Tree/Transformation/Base.hs b/src/Juvix/Compiler/Tree/Transformation/Base.hs
@@ -2,9 +2,11 @@ module Juvix.Compiler.Tree.Transformation.Base
   ( module Juvix.Compiler.Tree.Transformation.Generic.Base,
     module Juvix.Compiler.Tree.Language,
     module Juvix.Compiler.Tree.Data.InfoTable,
+    module Juvix.Compiler.Tree.Options,
   )
 where
 
 import Juvix.Compiler.Tree.Data.InfoTable
 import Juvix.Compiler.Tree.Language
+import Juvix.Compiler.Tree.Options
 import Juvix.Compiler.Tree.Transformation.Generic.Base
diff --git a/src/Juvix/Compiler/Tree/Transformation/Optimize/ConvertUnaryCalls.hs b/src/Juvix/Compiler/Tree/Transformation/Optimize/ConvertUnaryCalls.hs
@@ -0,0 +1,53 @@
+module Juvix.Compiler.Tree.Transformation.Optimize.ConvertUnaryCalls where
+
+import Juvix.Compiler.Core.Data.BinderList qualified as BL
+import Juvix.Compiler.Tree.Data.InfoTable
+import Juvix.Compiler.Tree.Extra.Recursors
+import Juvix.Compiler.Tree.Extra.Type
+import Juvix.Compiler.Tree.Transformation.Base
+
+-- | Replaces generic calls (with CallClosures) to unknown unary functions with
+-- known non-function target types by direct closure calls (with Call)
+convertUnaryCalls :: InfoTable -> InfoTable
+convertUnaryCalls tab = mapT convert tab
+  where
+    convert :: Symbol -> Node -> Node
+    convert sym = umapL (go argtys)
+      where
+        funInfo = lookupFunInfo tab sym
+        argtys
+          | funInfo ^. functionArgsNum == 0 = []
+          | otherwise = typeArgs (funInfo ^. functionType)
+
+    go :: [Type] -> BinderList TempVar -> Node -> Node
+    go argtys tmps node = case node of
+      CallClosures ncl@NodeCallClosures {..}
+        | length _nodeCallClosuresArgs == 1 ->
+            case _nodeCallClosuresFun of
+              MemRef NodeMemRef {..}
+                | DRef (ArgRef OffsetRef {..}) <- _nodeMemRef,
+                  isUnaryWithAtomicTarget (argtys !! _offsetRefOffset) ->
+                    mkClosureCall ncl
+                | DRef (TempRef (RefTemp OffsetRef {..})) <- _nodeMemRef,
+                  isUnaryWithAtomicTarget (BL.lookupLevel _offsetRefOffset tmps ^. tempVarType) ->
+                    mkClosureCall ncl
+                | ConstrRef (Field {..}) <- _nodeMemRef,
+                  constrInfo <- lookupConstrInfo tab _fieldTag,
+                  isUnaryWithAtomicTarget (typeArgs (constrInfo ^. constructorType) !! _fieldOffset) ->
+                    mkClosureCall ncl
+              _ -> node
+      _ -> node
+
+    mkClosureCall :: NodeCallClosures -> Node
+    mkClosureCall NodeCallClosures {..} =
+      Call
+        NodeCall
+          { _nodeCallInfo = _nodeCallClosuresInfo,
+            _nodeCallType = CallClosure _nodeCallClosuresFun,
+            _nodeCallArgs = toList _nodeCallClosuresArgs
+          }
+
+    isUnaryWithAtomicTarget :: Type -> Bool
+    isUnaryWithAtomicTarget ty =
+      length (typeArgs ty) == 1
+        && isConcreteAtomType (typeTarget ty)
diff --git a/src/Juvix/Compiler/Tree/Transformation/Optimize/Phase/Main.hs b/src/Juvix/Compiler/Tree/Transformation/Optimize/Phase/Main.hs
@@ -0,0 +1,9 @@
+module Juvix.Compiler.Tree.Transformation.Optimize.Phase.Main where
+
+import Juvix.Compiler.Tree.Transformation.Base
+import Juvix.Compiler.Tree.Transformation.Optimize.ConvertUnaryCalls
+
+optimize :: (Member (Reader Options) r) => InfoTable -> Sem r InfoTable
+optimize =
+  withOptimizationLevel 1 $
+    return . convertUnaryCalls
diff --git a/src/Juvix/Compiler/Tree/Translation/FromAsm.hs b/src/Juvix/Compiler/Tree/Translation/FromAsm.hs
@@ -192,7 +192,12 @@ goFunction infoTab fi = do
             Save
               NodeSave
                 { _nodeSaveInfo = mempty,
-                  _nodeSaveTempVar = TempVar _cmdSaveName (_cmdSaveInfo ^. Asm.commandInfoLocation),
+                  _nodeSaveTempVar =
+                    TempVar
+                      { _tempVarName = _cmdSaveName,
+                        _tempVarLocation = _cmdSaveInfo ^. Asm.commandInfoLocation,
+                        _tempVarType = TyDynamic
+                      },
                   _nodeSaveArg = arg,
                   _nodeSaveBody = body
                 }
@@ -255,7 +260,7 @@ goFunction infoTab fi = do
               NodeSave
                 { _nodeSaveInfo = mempty,
                   _nodeSaveArg = arg,
-                  _nodeSaveTempVar = TempVar Nothing Nothing,
+                  _nodeSaveTempVar = TempVar Nothing Nothing TyDynamic,
                   _nodeSaveBody =
                     Binop
                       NodeBinop

diff --git a/src/Juvix/Compiler/Tree/Translation/FromCore.hs b/src/Juvix/Compiler/Tree/Translation/FromCore.hs
@@ -270,7 +270,12 @@ genCode infoTable fi =
           { _nodeSaveInfo = mempty,
             _nodeSaveArg = arg,
             _nodeSaveBody = body,
-            _nodeSaveTempVar = TempVar (Just name) loc
+            _nodeSaveTempVar =
+              TempVar
+                { _tempVarName = Just name,
+                  _tempVarLocation = loc,
+                  _tempVarType = convertType 0 (_letItem ^. Core.letItemBinder . Core.binderType)
+                }
           }
       where
         name = _letItem ^. Core.letItemBinder . Core.binderName

diff --git a/src/Juvix/Compiler/Tree/Translation/FromSource.hs b/src/Juvix/Compiler/Tree/Translation/FromSource.hs
@@ -430,7 +430,12 @@ parseSave = do
       { _nodeSaveInfo = NodeInfo (Just loc'),
         _nodeSaveArg = arg,
         _nodeSaveBody = body,
-        _nodeSaveTempVar = TempVar mname (Just loc)
+        _nodeSaveTempVar =
+          TempVar
+            { _tempVarName = mname,
+              _tempVarLocation = Just loc,
+              _tempVarType = TyDynamic
+            }
       }
 
 withSave ::

diff --git a/test/Tree/Asm/Base.hs b/test/Tree/Asm/Base.hs
@@ -8,18 +8,20 @@ import Juvix.Compiler.Tree.Translation.FromSource
 import Juvix.Data.PPOutput
 
 treeAsmAssertion ::
+  Path Abs Dir ->
   Path Abs File ->
   Path Abs File ->
   (String -> IO ()) ->
   Assertion
-treeAsmAssertion mainFile expectedFile step = do
+treeAsmAssertion root' mainFile expectedFile step = do
   step "Parse"
   s <- readFile mainFile
   case runParser mainFile s of
     Left err -> assertFailure (prettyString err)
     Right tabIni -> do
       step "Translate"
-      case run $ runError @JuvixError $ Tree.toAsm tabIni of
+      entryPoint <- testDefaultEntryPointIO root' mainFile
+      case run $ runReader entryPoint $ runError @JuvixError $ Tree.toAsm tabIni of
         Left err -> assertFailure (prettyString (fromJuvixError @GenericError err))
         Right tab -> do
           let tab' = Asm.fromTree tab

diff --git a/test/Tree/Asm/Positive.hs b/test/Tree/Asm/Positive.hs
@@ -12,7 +12,7 @@ testDescr Eval.PosTest {..} =
    in TestDescr
         { _testName = _name,
           _testRoot = tRoot,
-          _testAssertion = Steps $ treeAsmAssertion file' expected'
+          _testAssertion = Steps $ treeAsmAssertion tRoot file' expected'
         }
 
 ignoredTests :: [String]

diff --git a/test/Tree/Eval/Base.hs b/test/Tree/Eval/Base.hs
@@ -7,6 +7,7 @@ import Juvix.Compiler.Tree.Error
 import Juvix.Compiler.Tree.Evaluator
 import Juvix.Compiler.Tree.Language.Base
 import Juvix.Compiler.Tree.Language.Value
+import Juvix.Compiler.Tree.Options qualified as Tree
 import Juvix.Compiler.Tree.Pretty
 import Juvix.Compiler.Tree.Transformation
 import Juvix.Compiler.Tree.Translation.FromSource
@@ -36,12 +37,13 @@ treeEvalAssertionParam evalParam mainFile expectedFile trans testTrans step = do
     Left err -> assertFailure (prettyString err)
     Right tab0 -> do
       step "Validate"
-      case run $ runError @JuvixError $ applyTransformations [Validate] tab0 of
+      let opts = Tree.defaultOptions
+      case run $ runReader opts $ runError @JuvixError $ applyTransformations [Validate] tab0 of
         Left err -> assertFailure (prettyString (fromJuvixError @GenericError err))
         Right tab1 -> do
           unless (null trans) $
             step "Transform"
-          case run $ runError @JuvixError $ applyTransformations trans tab1 of
+          case run $ runReader opts $ runError @JuvixError $ applyTransformations trans tab1 of
             Left err -> assertFailure (prettyString (fromJuvixError @GenericError err))
             Right tab -> do
               testTrans tab

diff --git a/test/Tree/Eval/Positive.hs b/test/Tree/Eval/Positive.hs
@@ -249,5 +249,10 @@ tests =
       "Test042: Uncurried function type unification"
       $(mkRelDir ".")
       $(mkRelFile "test042.jvt")
-      $(mkRelFile "out/test042.out")
+      $(mkRelFile "out/test042.out"),
+    PosTest
+      "Test043: Unary closure calls"
+      $(mkRelDir ".")
+      $(mkRelFile "test043.jvt")
+      $(mkRelFile "out/test043.out")
   ]