Add hook for custom analysis (#99)

backend-es/src/Main.purs

@@ -250,6 +250,7 @@ main cliRoot =
   buildCmd args = liftEffect makeBuildState >>= \state -> basicBuildMain
     { resolveCoreFnDirectory: pure args.coreFnDir
     , resolveExternalDirectives: map (fromMaybe Map.empty) $ traverse externalDirectivesFromFile args.directivesFile
+    , analyzeCustom: \_ _ -> Nothing
     , foreignSemantics: Map.union coreForeignSemantics esForeignSemantics
     , onCodegenBefore: do
         liftEffect $ flip Ref.write state.codegenStartTime <<< Just =<< now

backend-es/src/PureScript/Backend/Optimizer/Codegen/EcmaScript/Builder.purs

@@ -15,7 +15,7 @@ import Data.Lazy as Lazy
 import Data.List (List)
 import Data.Map (Map)
 import Data.Map as Map
-import Data.Maybe (maybe)
+import Data.Maybe (Maybe, maybe)
 import Data.Set (Set)
 import Data.Set as Set
 import Data.Set.NonEmpty as NonEmptySet
@@ -28,6 +28,7 @@ import Node.FS.Aff as FS
 import Node.Glob.Basic (expandGlobs)
 import Node.Path (FilePath)
 import Node.Process as Process
+import PureScript.Backend.Optimizer.Analysis (BackendAnalysis)
 import PureScript.Backend.Optimizer.Builder (BuildEnv, buildModules)
 import PureScript.Backend.Optimizer.Convert (BackendModule, OptimizationSteps)
 import PureScript.Backend.Optimizer.CoreFn (Ann, Ident, Module, ModuleName(..), Qualified)
@@ -35,8 +36,9 @@ import PureScript.Backend.Optimizer.CoreFn.Json (decodeModule)
 import PureScript.Backend.Optimizer.CoreFn.Sort (emptyPull, pullResult, resumePull, sortModules)
 import PureScript.Backend.Optimizer.Directives (parseDirectiveFile)
 import PureScript.Backend.Optimizer.Directives.Defaults as Defaults
-import PureScript.Backend.Optimizer.Semantics (InlineDirectiveMap)
+import PureScript.Backend.Optimizer.Semantics (BackendExpr, Ctx, InlineDirectiveMap)
 import PureScript.Backend.Optimizer.Semantics.Foreign (ForeignEval)
+import PureScript.Backend.Optimizer.Syntax (BackendSyntax)
 import PureScript.CST.Errors (printParseError)
 
 coreFnModulesFromOutput :: String -> NonEmptyArray String -> Aff (Either (NonEmptyArray (Tuple FilePath String)) (List (Module Ann)))
@@ -83,6 +85,7 @@ externalDirectivesFromFile filePath = do
 basicBuildMain
   :: { resolveCoreFnDirectory :: Aff FilePath
      , resolveExternalDirectives :: Aff InlineDirectiveMap
+     , analyzeCustom :: Ctx -> BackendSyntax BackendExpr -> Maybe BackendAnalysis
      , foreignSemantics :: Map (Qualified Ident) ForeignEval
      , onCodegenBefore :: Aff Unit
      , onCodegenAfter :: Aff Unit
@@ -106,7 +109,8 @@ basicBuildMain options = do
     Right coreFnModules -> do
       options.onCodegenBefore
       coreFnModules # buildModules
-        { directives: allDirectives
+        { analyzeCustom: options.analyzeCustom
+        , directives: allDirectives
         , foreignSemantics: options.foreignSemantics
         , onCodegenModule: options.onCodegenModule
         , onPrepareModule: options.onPrepareModule

backend-es/test/Main.purs

@@ -120,6 +120,7 @@ runSnapshotTests { accept, filter, traceIdents } = do
       stepsRef <- liftEffect $ Ref.new []
       coreFnModules # buildModules
         { directives
+        , analyzeCustom: \_ _ -> Nothing
         , foreignSemantics: Map.union coreForeignSemantics esForeignSemantics
         , onCodegenModule: \build (Module { name: ModuleName name, path }) backendMod optimizationSteps -> do
             let

src/PureScript/Backend/Optimizer/Analysis.purs

@@ -200,14 +200,19 @@ class HasAnalysis a where
 resultOf :: forall a. HasAnalysis a => a -> ResultTerm
 resultOf = analysisOf >>> unwrap >>> _.result
 
-analyze :: forall a. HasAnalysis a => HasSyntax a => (Tuple (Qualified Ident) (Maybe BackendAccessor) -> BackendAnalysis) -> BackendSyntax a -> BackendAnalysis
+analyze :: forall a. HasAnalysis a => HasSyntax a => (Qualified Ident -> Maybe String -> Maybe BackendAnalysis) -> BackendSyntax a -> BackendAnalysis
 analyze externAnalysis expr = case expr of
-  Var qi -> do
-    let BackendAnalysis { args } = externAnalysis (Tuple qi Nothing)
-    withArgs args
-      $ bump
-      $ externs
-      $ usedDep qi mempty
+  Var qi ->
+    case externAnalysis qi Nothing of
+      Just (BackendAnalysis { args }) ->
+        withArgs args analysis
+      Nothing ->
+        analysis
+    where
+    analysis =
+      bump
+        $ externs
+        $ usedDep qi mempty
   Local _ lvl ->
     bump
       $ used lvl
@@ -351,9 +356,13 @@ analyze externAnalysis expr = case expr of
         analysis
       Just (Local _ lvl) ->
         accessed lvl $ complex Deref analysis
-      Just (Var qi) -> do
-        let BackendAnalysis { args } = externAnalysis (Tuple qi (Just acc))
-        withArgs args $ complex Trivial analysis
+      Just (Var qi) ->
+        case acc of
+          GetProp prop
+            | Just (BackendAnalysis { args }) <- externAnalysis qi (Just prop) ->
+                withArgs args $ complex Trivial analysis
+          _ ->
+            complex Trivial analysis
       _ ->
         complex Deref analysis
     where
@@ -381,7 +390,7 @@ analyze externAnalysis expr = case expr of
       withResult KnownNeutral
         $ analyzeDefault expr
 
-analyzeEffectBlock :: forall a. HasAnalysis a => HasSyntax a => (Tuple (Qualified Ident) (Maybe BackendAccessor) -> BackendAnalysis) -> BackendSyntax a -> BackendAnalysis
+analyzeEffectBlock :: forall a. HasAnalysis a => HasSyntax a => (Qualified Ident -> Maybe String -> Maybe BackendAnalysis) -> BackendSyntax a -> BackendAnalysis
 analyzeEffectBlock externAnalysis expr = case expr of
   Let _ lvl a b ->
     withResult (resultOf b)

src/PureScript/Backend/Optimizer/Builder.purs

@@ -11,13 +11,15 @@ import Data.List (List, foldM)
 import Data.List as List
 import Data.Map (Map)
 import Data.Map as Map
+import Data.Maybe (Maybe)
 import Data.Set (Set)
 import Data.Tuple (Tuple(..))
 import PureScript.Backend.Optimizer.Analysis (BackendAnalysis)
 import PureScript.Backend.Optimizer.Convert (BackendModule, OptimizationSteps, toBackendModule)
 import PureScript.Backend.Optimizer.CoreFn (Ann, Ident, Module(..), Qualified)
-import PureScript.Backend.Optimizer.Semantics (ExternImpl, InlineDirectiveMap)
+import PureScript.Backend.Optimizer.Semantics (BackendExpr, Ctx, ExternImpl, InlineDirectiveMap)
 import PureScript.Backend.Optimizer.Semantics.Foreign (ForeignEval)
+import PureScript.Backend.Optimizer.Syntax (BackendSyntax)
 
 type BuildEnv =
   { implementations :: Map (Qualified Ident) (Tuple BackendAnalysis ExternImpl)
@@ -26,7 +28,8 @@ type BuildEnv =
   }
 
 type BuildOptions m =
-  { directives :: InlineDirectiveMap
+  { analyzeCustom :: Ctx -> BackendSyntax BackendExpr -> Maybe BackendAnalysis
+  , directives :: InlineDirectiveMap
   , foreignSemantics :: Map (Qualified Ident) ForeignEval
   , onPrepareModule :: BuildEnv -> Module Ann -> m (Module Ann)
   , onCodegenModule :: BuildEnv -> Module Ann -> BackendModule -> OptimizationSteps -> m Unit
@@ -45,7 +48,8 @@ buildModules options coreFnModules =
     coreFnModule'@(Module { name }) <- options.onPrepareModule buildEnv coreFnModule
     let
       Tuple optimizationSteps backendMod = toBackendModule coreFnModule'
-        { currentModule: name
+        { analyzeCustom: options.analyzeCustom
+        , currentModule: name
         , currentLevel: 0
         , toLevel: Map.empty
         , implementations

src/PureScript/Backend/Optimizer/Convert.purs

@@ -70,10 +70,10 @@ import Data.Traversable (class Foldable, Accum, foldr, for, mapAccumL, mapAccumR
 import Data.TraversableWithIndex (forWithIndex)
 import Data.Tuple (Tuple(..), fst, snd)
 import Partial.Unsafe (unsafeCrashWith, unsafePartial)
-import PureScript.Backend.Optimizer.Analysis (BackendAnalysis)
+import PureScript.Backend.Optimizer.Analysis (BackendAnalysis, analyze, analyzeEffectBlock)
 import PureScript.Backend.Optimizer.CoreFn (Ann(..), Bind(..), Binder(..), Binding(..), CaseAlternative(..), CaseGuard(..), Comment, ConstructorType(..), Expr(..), Guard(..), Ident(..), Literal(..), Meta(..), Module(..), ModuleName(..), ProperName, Qualified(..), ReExport, findProp, propKey, propValue, qualifiedModuleName, unQualified)
 import PureScript.Backend.Optimizer.Directives (DirectiveHeaderResult, parseDirectiveHeader)
-import PureScript.Backend.Optimizer.Semantics (BackendExpr(..), BackendSemantics, Ctx, DataTypeMeta, Env(..), EvalRef(..), ExternImpl(..), ExternSpine, InlineAccessor(..), InlineDirective(..), InlineDirectiveMap, NeutralExpr(..), build, evalExternFromImpl, evalExternRefFromImpl, freeze, optimize)
+import PureScript.Backend.Optimizer.Semantics (BackendExpr(..), BackendSemantics, Ctx(..), DataTypeMeta, Env(..), EvalRef(..), ExternImpl(..), ExternSpine, InlineAccessor(..), InlineDirective(..), InlineDirectiveMap, NeutralExpr(..), build, evalExternFromImpl, evalExternRefFromImpl, freeze, optimize)
 import PureScript.Backend.Optimizer.Semantics.Foreign (ForeignEval)
 import PureScript.Backend.Optimizer.Syntax (BackendAccessor(..), BackendOperator(..), BackendOperator1(..), BackendOperator2(..), BackendOperatorOrd(..), BackendSyntax(..), Level(..), Pair(..))
 import PureScript.Backend.Optimizer.Utils (foldl1Array)
@@ -100,7 +100,8 @@ type BackendModule =
   }
 
 type ConvertEnv =
-  { currentLevel :: Int
+  { analyzeCustom :: Ctx -> BackendSyntax BackendExpr -> Maybe BackendAnalysis
+  , currentLevel :: Int
   , currentModule :: ModuleName
   , dataTypes :: Map ProperName DataTypeMeta
   , toLevel :: Map Ident Level
@@ -357,29 +358,33 @@ buildM :: BackendSyntax BackendExpr -> ConvertM BackendExpr
 buildM a env = build (getCtx env) a
 
 getCtx :: ConvertEnv -> Ctx
-getCtx env =
+getCtx env = Ctx
   { currentLevel: env.currentLevel
   , lookupExtern
+  , analyze: \ctx@(Ctx { effect }) expr ->
+      case env.analyzeCustom ctx expr of
+        Just s -> s
+        Nothing ->
+          if effect then
+            analyzeEffectBlock lookupExtern expr
+          else
+            analyze lookupExtern expr
   , effect: false
   }
   where
-  lookupExtern (Tuple qual acc) = do
+  lookupExtern qual acc = do
     Tuple s impl <- Map.lookup qual env.implementations
     case impl of
-      ExternExpr _ a ->
+      ExternExpr _ _ ->
         case acc of
-          Nothing ->
-            Just (Tuple s a)
-          _ ->
-            Nothing
+          Nothing -> Just s
+          _ -> Nothing
       ExternDict _ a ->
         case acc of
-          Just (GetProp prop) ->
-            findProp prop a
-          -- Nothing ->
-          --   Just $ Tuple s $ NeutralExpr $ Lit $ LitRecord (map snd <$> a)
-          _ ->
-            Nothing
+          Just prop ->
+            fst <$> findProp prop a
+          Nothing ->
+            Just s
       ExternCtor _ _ _ _ _ ->
         Nothing
 

src/PureScript/Backend/Optimizer/Semantics.purs

@@ -22,7 +22,7 @@ import Data.Set as Set
 import Data.String as String
 import Data.Tuple (Tuple(..), fst, snd)
 import Partial.Unsafe (unsafeCrashWith)
-import PureScript.Backend.Optimizer.Analysis (class HasAnalysis, BackendAnalysis(..), Capture(..), Complexity(..), ResultTerm(..), Usage(..), analysisOf, analyze, analyzeEffectBlock, bound, bump, complex, resultOf, updated, withResult, withRewrite)
+import PureScript.Backend.Optimizer.Analysis (class HasAnalysis, BackendAnalysis(..), Capture(..), Complexity(..), ResultTerm(..), Usage(..), analysisOf, bound, bump, complex, resultOf, updated, withResult, withRewrite)
 import PureScript.Backend.Optimizer.CoreFn (ConstructorType, Ident(..), Literal(..), ModuleName, Prop(..), ProperName, Qualified(..), findProp, propKey, propValue)
 import PureScript.Backend.Optimizer.Syntax (class HasSyntax, BackendAccessor(..), BackendEffect, BackendOperator(..), BackendOperator1(..), BackendOperator2(..), BackendOperatorNum(..), BackendOperatorOrd(..), BackendSyntax(..), Level(..), Pair(..), syntaxOf)
 import PureScript.Backend.Optimizer.Utils (foldl1Array, foldr1Array)
@@ -1150,14 +1150,25 @@ caseNumber = case _ of
   NeutLit (LitNumber a) -> Just a
   _ -> Nothing
 
-type Ctx =
+newtype Ctx = Ctx
   { currentLevel :: Int
-  , lookupExtern :: Tuple (Qualified Ident) (Maybe BackendAccessor) -> Maybe (Tuple BackendAnalysis NeutralExpr)
+  , lookupExtern :: Qualified Ident -> Maybe String -> Maybe BackendAnalysis
+  , analyze :: Ctx -> BackendSyntax BackendExpr -> BackendAnalysis
   , effect :: Boolean
   }
 
 nextLevel :: Ctx -> Tuple Level Ctx
-nextLevel ctx = Tuple (Level ctx.currentLevel) $ ctx { currentLevel = ctx.currentLevel + 1 }
+nextLevel (Ctx ctx) = Tuple (Level ctx.currentLevel) $ Ctx ctx { currentLevel = ctx.currentLevel + 1 }
+
+effectfully :: Ctx -> Ctx
+effectfully (Ctx ctx)
+  | ctx.effect = Ctx ctx
+  | otherwise = Ctx ctx { effect = true }
+
+purely :: Ctx -> Ctx
+purely (Ctx ctx)
+  | ctx.effect = Ctx ctx { effect = false }
+  | otherwise = Ctx ctx
 
 quote :: Ctx -> BackendSemantics -> BackendExpr
 quote = go
@@ -1166,7 +1177,7 @@ quote = go
     -- Block constructors
     SemLet ident binding k -> do
       let Tuple level ctx' = nextLevel ctx
-      build ctx $ Let ident level (quote (ctx { effect = false }) binding) $ quote ctx' $ k $ SemRef (EvalLocal ident level) [] $ defer \_ -> deref binding
+      build ctx $ Let ident level (quote (purely ctx) binding) $ quote ctx' $ k $ SemRef (EvalLocal ident level) [] $ defer \_ -> deref binding
     SemLetRec bindings k -> do
       let Tuple level ctx' = nextLevel ctx
       -- We are not currently propagating references
@@ -1175,18 +1186,18 @@ quote = go
       -- which we don't currently implement.
       let neutBindings = (\(Tuple ident _) -> Tuple ident $ defer \_ -> NeutLocal (Just ident) level) <$> bindings
       build ctx $ LetRec level
-        (map (\b -> quote (ctx' { effect = false }) $ b neutBindings) <$> bindings)
+        (map (\b -> quote (purely ctx') $ b neutBindings) <$> bindings)
         (quote ctx' $ k neutBindings)
     SemEffectBind ident binding k -> do
-      let ctx' = ctx { effect = true }
+      let ctx' = effectfully ctx
       let Tuple level ctx'' = nextLevel ctx'
       build ctx $ EffectBind ident level (quote ctx' binding) $ quote ctx'' $ k $ NeutLocal ident level
     SemEffectPure sem ->
-      build ctx $ EffectPure (quote (ctx { effect = false }) sem)
+      build ctx $ EffectPure (quote (purely ctx) sem)
     SemEffectDefer sem ->
-      build ctx $ EffectDefer (quote (ctx { effect = true }) sem)
+      build ctx $ EffectDefer (quote (effectfully ctx) sem)
     SemBranch branches def -> do
-      let ctx' = ctx { effect = false }
+      let ctx' = purely ctx
       let quoteCond (SemConditional a b) = Pair (quote ctx' $ force a) (quote ctx $ force b)
       let branches' = quoteCond <$> branches
       foldr (buildBranchCond ctx) (quote ctx <<< force $ def) branches'
@@ -1200,7 +1211,7 @@ quote = go
           go ctx $ neutralSpine (NeutLocal ident lvl) sp
     SemLam ident k -> do
       let Tuple level ctx' = nextLevel ctx
-      build ctx $ Abs (NonEmptyArray.singleton (Tuple ident level)) $ quote (ctx' { effect = false }) $ k $ NeutLocal ident level
+      build ctx $ Abs (NonEmptyArray.singleton (Tuple ident level)) $ quote (purely ctx') $ k $ NeutLocal ident level
     SemMkFn pro -> do
       let
         loop ctx' idents = case _ of
@@ -1208,7 +1219,7 @@ quote = go
             let Tuple lvl ctx'' = nextLevel ctx'
             loop ctx'' (Array.snoc idents (Tuple ident lvl)) (k (NeutLocal ident lvl))
           MkFnApplied body ->
-            build ctx' $ UncurriedAbs idents $ quote (ctx' { effect = false }) body
+            build ctx' $ UncurriedAbs idents $ quote (purely ctx') body
       loop ctx [] pro
     SemMkEffectFn pro -> do
       let
@@ -1217,7 +1228,7 @@ quote = go
             let Tuple lvl ctx'' = nextLevel ctx'
             loop ctx'' (Array.snoc idents (Tuple ident lvl)) (k (NeutLocal ident lvl))
           MkFnApplied body ->
-            build ctx' $ UncurriedEffectAbs idents $ quote (ctx' { effect = false }) body
+            build ctx' $ UncurriedEffectAbs idents $ quote (purely ctx') body
       loop ctx [] pro
     SemAssocOp op spine ->
       foldl1Array
@@ -1240,15 +1251,15 @@ quote = go
     NeutCtorDef _ ct ty tag fields ->
       build ctx $ CtorDef ct ty tag fields
     NeutUncurriedApp hd spine -> do
-      let ctx' = ctx { effect = false }
+      let ctx' = purely ctx
       let hd' = quote ctx' hd
       build ctx $ UncurriedApp hd' (quote ctx' <$> spine)
     NeutUncurriedEffectApp hd spine -> do
-      let ctx' = ctx { effect = false }
+      let ctx' = purely ctx
       let hd' = quote ctx' hd
       build ctx $ UncurriedEffectApp hd' (quote ctx' <$> spine)
     NeutApp hd spine -> do
-      let ctx' = ctx { effect = false }
+      let ctx' = purely ctx
       let hd' = quote ctx' hd
       case NonEmptyArray.fromArray (quote ctx' <$> spine) of
         Nothing ->
@@ -1264,7 +1275,7 @@ quote = go
     NeutPrimOp op ->
       build ctx $ PrimOp (quote ctx <$> op)
     NeutPrimEffect eff ->
-      build ctx $ PrimEffect (quote (ctx { effect = false }) <$> eff)
+      build ctx $ PrimEffect (quote (purely ctx) <$> eff)
     NeutPrimUndefined ->
       build ctx PrimUndefined
     NeutFail err ->
@@ -1397,13 +1408,10 @@ simplifyCondLiftAnd ctx pair def1 = case pair of
     Nothing
 
 buildStop :: Ctx -> Qualified Ident -> BackendExpr
-buildStop ctx stop = ExprRewrite (analyzeDefault ctx (Var stop)) (RewriteStop stop)
+buildStop ctx@(Ctx { analyze }) stop = ExprRewrite (analyze ctx (Var stop)) (RewriteStop stop)
 
 buildDefault :: Ctx -> BackendSyntax BackendExpr -> BackendExpr
-buildDefault ctx expr = ExprSyntax (analyzeDefault ctx expr) expr
-
-analyzeDefault :: Ctx -> BackendSyntax BackendExpr -> BackendAnalysis
-analyzeDefault ctx = (if ctx.effect then analyzeEffectBlock else analyze) (foldMap fst <<< ctx.lookupExtern)
+buildDefault ctx@(Ctx { analyze }) expr = ExprSyntax (analyze ctx expr) expr
 
 rewriteInline :: Maybe Ident -> Level -> BackendExpr -> BackendExpr -> BackendExpr
 rewriteInline ident level binding body = do