Annotation-based inlining + a few optimisations.

Author: Unisay

lib/Language/PureScript/Backend.hs

@@ -13,6 +13,7 @@ import Language.PureScript.Backend.Lua.Types qualified as Lua
 import Language.PureScript.Backend.Types (AppOrModule (..), entryPointModule)
 import Language.PureScript.CoreFn.Reader qualified as CoreFn
 import Path (Abs, Dir, Path, SomeBase)
+import Text.Pretty.Simple (pPrint)
 import Prelude hiding (show)
 
 data CompilationResult = CompilationResult
@@ -32,8 +33,8 @@ compileModules
   → AppOrModule
   → ExceptT (Variant e) IO CompilationResult
 compileModules outputDir foreignDir appOrModule = do
-  cfnModules ←
-    CoreFn.readModuleRecursively outputDir (entryPointModule appOrModule)
+  let entryModuleName = entryPointModule appOrModule
+  cfnModules ← CoreFn.readModuleRecursively outputDir entryModuleName
   let dataDecls = IR.collectDataDeclarations cfnModules
   irResults ← forM (Map.toList cfnModules) \(_psModuleName, cfnModule) →
     Oops.hoistEither $ IR.mkModule cfnModule dataDecls
@@ -42,12 +43,8 @@ compileModules outputDir foreignDir appOrModule = do
         Linker.makeUberModule (linkerMode appOrModule) irModules
           & optimizedUberModule
   let needsRuntimeLazy = Tagged (any untag needsRuntimeLazys)
-
-  unoptimizedChunk ←
-    Lua.fromUberModule foreignDir needsRuntimeLazy appOrModule uberModule
-  pure
-    CompilationResult {lua = optimizeChunk unoptimizedChunk, ir = uberModule}
-
+  chunk ← Lua.fromUberModule foreignDir needsRuntimeLazy appOrModule uberModule
+  pure CompilationResult {lua = optimizeChunk chunk, ir = uberModule}
 linkerMode ∷ AppOrModule → Linker.LinkMode
 linkerMode = \case
   AsModule psModuleName → Linker.LinkAsModule psModuleName

lib/Language/PureScript/Backend/IR.hs

@@ -69,9 +69,13 @@ runRepM
   ∷ Context
   → RepM a
   → Either CoreFnError (Tagged "needsRuntimeLazy" Bool, a)
-runRepM ctx (RepM m) =
-  runStateT m ctx <&> \(a, ctx') →
-    (Tagged . getAny $ needsRuntimeLazy ctx', a)
+runRepM ctx (RepM m) = do
+  (a, ctx') ← runStateT m ctx
+  let remainingAnnotations = annotations ctx'
+  unless (Map.null remainingAnnotations) do
+    Left . CoreFnError (Cfn.moduleName (contextModule ctx)) $
+      UnusedAnnotations remainingAnnotations
+  pure (Tagged . getAny $ needsRuntimeLazy ctx', a)
 
 mkModule
   ∷ Cfn.Module Cfn.Ann
@@ -119,6 +123,15 @@ parseAnnotations currentModule =
       & first
         (CoreFnError (Cfn.moduleName currentModule) . AnnotationParsingError)
 
+useAnnotation ∷ Name → RepM (Maybe Annotation)
+useAnnotation name = do
+  ctx ← get
+  let (ann, annotations') =
+        -- delete the annotation from the map returning the value
+        Map.updateLookupWithKey (\_ _ → Nothing) name (annotations ctx)
+  put $ ctx {annotations = annotations'}
+  pure ann
+
 mkImports ∷ RepM [ModuleName]
 mkImports = do
   Cfn.Module {moduleName, moduleImports} ← gets contextModule
@@ -135,8 +148,13 @@ mkReExports =
   Map.fromAscList . fmap (identToName <<$>>) . Map.toAscList
     <$> gets (contextModule >>> Cfn.moduleReExports)
 
-mkForeign ∷ RepM [Name]
-mkForeign = identToName <<$>> gets (contextModule >>> Cfn.moduleForeign)
+mkForeign ∷ RepM [(Ann, Name)]
+mkForeign = do
+  idents ← gets (contextModule >>> Cfn.moduleForeign)
+  forM idents \ident → do
+    let name = identToName ident
+    ann ← useAnnotation name
+    pure (ann, name)
 
 collectDataDeclarations
   ∷ Map ModuleName (Cfn.Module Cfn.Ann)
@@ -179,7 +197,7 @@ mkBinding ∷ Cfn.Bind Cfn.Ann → RepM Binding
 mkBinding = \case
   Cfn.NonRec _ann ident cfnExpr → do
     let name = identToName ident
-    ann ← gets $ annotations >>> Map.lookup name
+    ann ← useAnnotation name
     expr ← makeExprAnnotated ann cfnExpr
     pure $ Standalone (noAnn, name, expr)
   Cfn.Rec bindingGroup → do
@@ -324,7 +342,7 @@ mkLet ann binds expr = do
 -- The algorithm is based on this document: ------------------------------------
 -- https://julesjacobs.com/notes/patternmatching/patternmatching.pdf -----------
 
-mkCase ∷ Ann -> [CfnExp] → NonEmpty (Cfn.CaseAlternative Cfn.Ann) → RepM Exp
+mkCase ∷ Ann → [CfnExp] → NonEmpty (Cfn.CaseAlternative Cfn.Ann) → RepM Exp
 mkCase ann cfnExpressions alternatives = do
   expressions ← traverse makeExpr cfnExpressions
   -- Before making clauses, we need to prepare bindings
@@ -699,7 +717,10 @@ algebraicTy modName tyName = do
 --------------------------------------------------------------------------------
 -- Errors ----------------------------------------------------------------------
 
-throwContextualError ∷ CoreFnErrorReason → RepM a
+throwContextualError
+  ∷ (MonadState Context m, MonadError CoreFnError m)
+  ⇒ CoreFnErrorReason
+  → m a
 throwContextualError e = do
   currentModule ← gets (contextModule >>> Cfn.moduleName)
   throwError $ CoreFnError currentModule e
@@ -730,6 +751,7 @@ data CoreFnErrorReason
       TyName
   | UnicodeDecodeError UnicodeException
   | AnnotationParsingError (Megaparsec.ParseErrorBundle Text Void)
+  | UnusedAnnotations (Map Name Annotation)
 
 instance Show CoreFnErrorReason where
   show = \case
@@ -762,3 +784,5 @@ instance Show CoreFnErrorReason where
       "Unicode decode error: " <> displayException e
     AnnotationParsingError bundle →
       "Annotation parsing error: " <> Megaparsec.errorBundlePretty bundle
+    UnusedAnnotations anns →
+      "Unused annotations: " <> toString (pShow anns)

lib/Language/PureScript/Backend/IR/Inliner.hs

@@ -1,20 +1,13 @@
 module Language.PureScript.Backend.IR.Inliner where
 
-import Control.Monad.Combinators (choice)
 import Language.PureScript.Backend.IR.Names (Name, nameParser)
 import Text.Megaparsec qualified as Megaparsec
 import Text.Megaparsec.Char qualified as MC
 import Text.Megaparsec.Char.Lexer qualified as ML
 
 type Pragma = (Name, Annotation)
 
-data Annotation = Annotation InlineScope InlineRecipe
-  deriving stock (Show, Eq, Ord)
-
-data InlineScope = InModule | Global
-  deriving stock (Show, Eq, Ord)
-
-data InlineRecipe = Default | Always | Never
+data Annotation = Always | Never
   deriving stock (Show, Eq, Ord)
 
 type Parser = Megaparsec.Parsec Void Text
@@ -25,18 +18,7 @@ pragmaParser = do
   (,) <$> (nameParser <* sc) <*> annotationParser
 
 annotationParser ∷ Parser Annotation
-annotationParser = Annotation <$> scopeParser <*> recipeParser
-
-recipeParser ∷ Parser InlineRecipe
-recipeParser =
-  choice
-    [ Default <$ symbol "default"
-    , Always <$ symbol "always"
-    , Never <$ symbol "never"
-    ]
-
-scopeParser ∷ Parser InlineScope
-scopeParser = maybe InModule (const Global) <$> optional (symbol "export")
+annotationParser = (Always <$ symbol "always") <|> (Never <$ symbol "never")
 
 symbol ∷ Text → Parser ()
 symbol = void . ML.symbol sc

lib/Language/PureScript/Backend/IR/Linker.hs

@@ -4,6 +4,7 @@ module Language.PureScript.Backend.IR.Linker where
 
 import Data.Graph (graphFromEdges', reverseTopSort)
 import Data.Map qualified as Map
+import Language.PureScript.Backend.IR.Inliner qualified as Inline
 import Language.PureScript.Backend.IR.Names
   ( ModuleName
   , Name (..)
@@ -22,7 +23,6 @@ import Language.PureScript.Backend.IR.Types
   , RawExp (..)
   , bindingNames
   , noAnn
-  , objectProp
   , refImported
   )
 
@@ -81,10 +81,13 @@ foreignBindings Module {moduleName, modulePath, moduleForeigns} =
     ]
 
   foreignNamesBindings ∷ [Grouping (QName, Exp)] =
-    moduleForeigns <&> \name →
+    moduleForeigns <&> \(_ann, name) →
       Standalone
         ( QName moduleName name
-        , objectProp foreignModuleRef (PropName (nameToText name))
+        , ObjectProp
+            (Just Inline.Always)
+            foreignModuleRef
+            (PropName (nameToText name))
         )
 
 qualifiedModuleBindings ∷ Module → [Grouping (QName, Exp)]
@@ -98,7 +101,7 @@ qualifiedModuleBindings Module {moduleName, moduleBindings, moduleForeigns} =
     (QName moduleName name, qualifyTopRefs moduleName topRefs expr)
    where
     topRefs ∷ Map Name Index = Map.fromList do
-      (,0) <$> ((moduleBindings >>= bindingNames) <> moduleForeigns)
+      (,0) <$> ((moduleBindings >>= bindingNames) <> fmap snd moduleForeigns)
 
 qualifyTopRefs ∷ ModuleName → Map Name Index → Exp → Exp
 qualifyTopRefs moduleName = go

lib/Language/PureScript/Backend/IR/Optimizer.hs

@@ -4,11 +4,7 @@ import Data.List.NonEmpty qualified as NE
 import Data.Map qualified as Map
 import Data.Set qualified as Set
 import Language.PureScript.Backend.IR.DCE qualified as DCE
-import Language.PureScript.Backend.IR.Inliner
-  ( Annotation (..)
-  , InlineRecipe (..)
-  , InlineScope (InModule)
-  )
+import Language.PureScript.Backend.IR.Inliner (Annotation (..))
 import Language.PureScript.Backend.IR.Linker (UberModule (..))
 import Language.PureScript.Backend.IR.Names
   ( Name (..)
@@ -228,6 +224,7 @@ optimizedExpression =
   rewriteExpTopDown $
     constantFolding
       `thenRewrite` betaReduce
+      `thenRewrite` etaReduce
       `thenRewrite` betaReduceUnusedParams
       `thenRewrite` removeUnreachableThenBranch
       `thenRewrite` removeUnreachableElseBranch
@@ -239,6 +236,9 @@ constantFolding =
   pure . \case
     Eq _ (LiteralBool _ a) (LiteralBool _ b) →
       Rewritten Stop $ literalBool $ a == b
+    Eq _ (LiteralBool _ True) b →
+      -- 'b' must be of type Bool
+      Rewritten Stop b
     Eq _ (LiteralInt _ a) (LiteralInt _ b) →
       Rewritten Stop $ literalBool $ a == b
     Eq _ (LiteralFloat _ a) (LiteralFloat _ b) →
@@ -249,12 +249,20 @@ constantFolding =
       Rewritten Stop $ literalBool $ a == b
     _ → NoChange
 
--- \x -> (\y -> y + 1) x ===> (\y -> y + 1)
+-- (λx. M) N ===> M[x := N]
 betaReduce ∷ RewriteRule Ann
 betaReduce =
   pure . \case
-    Abs _ (ParamNamed _ _) (App _ f (Ref _ (Local _) 0)) →
-      Rewritten Recurse f
+    App _ (Abs _ (ParamNamed _ param) body) r →
+      Rewritten Recurse $ substitute (Local param) 0 r body
+    _ → NoChange
+
+-- (λx. M x) where x not free in M ===> M
+etaReduce ∷ RewriteRule Ann
+etaReduce =
+  pure . \case
+    Abs _ (ParamNamed _ _param) (App _ m (Ref _ (Local _) 0)) →
+      Rewritten Recurse m
     _ → NoChange
 
 betaReduceUnusedParams ∷ RewriteRule Ann
@@ -317,6 +325,6 @@ isInlinableExpr expr =
   hasInlineAnnotation ∷ Exp → Bool
   hasInlineAnnotation =
     getAnn >>> \case
-      Just (Annotation InModule Always) → True
-      Just (Annotation InModule Never) → False
-      _ → False
+      Just Always → True
+      Just Never → False
+      Nothing → False

lib/Language/PureScript/Backend/IR/Types.hs

@@ -33,7 +33,7 @@ data Module = Module
   , moduleImports ∷ [ModuleName]
   , moduleExports ∷ [Name]
   , moduleReExports ∷ Map ModuleName [Name]
-  , moduleForeigns ∷ [Name]
+  , moduleForeigns ∷ [(Ann, Name)]
   , modulePath ∷ FilePath
   }
 
@@ -83,7 +83,7 @@ data RawExp ann
   | Let ann (NonEmpty (Grouping (ann, Name, RawExp ann))) (RawExp ann)
   | IfThenElse ann (RawExp ann) (RawExp ann) (RawExp ann)
   | Exception ann Text
-  | ForeignImport ann ModuleName FilePath [Name]
+  | ForeignImport ann ModuleName FilePath [(ann, Name)]
 
 deriving stock instance Show ann ⇒ Show (RawExp ann)
 deriving stock instance Eq ann ⇒ Eq (RawExp ann)
@@ -359,7 +359,10 @@ annotateExpM around annotateExp annotateParam annotateName =
         pure $ IfThenElse ann i' t' e'
       Ctor _ann mn aty ty ctr fs → pure $ Ctor ann mn aty ty ctr fs
       Exception _ann m → pure $ Exception ann m
-      ForeignImport _ann m p ns → pure $ ForeignImport ann m p ns
+      ForeignImport _ann m p ns → do
+        anns ← traverse (uncurry annotateName) ns
+        let ns' = zip anns (fmap snd ns)
+        pure $ ForeignImport ann m p ns'
  where
   mkAnn ∷ RawExp ann → m (RawExp ann')
   mkAnn = annotateExpM around annotateExp annotateParam annotateName