Allow inlining of any function, do check dynamically

Author: jespercockx

src/Agda2Hs/Compile.hs

@@ -11,6 +11,7 @@ import Data.IORef
 import Data.List ( isPrefixOf, group, sort )
 
 import qualified Data.Map as M
+import qualified Data.Set as S
 
 import Agda.Compiler.Backend
 import Agda.Compiler.Common ( curIF )
@@ -41,7 +42,8 @@ globalSetup :: Options -> TCM GlobalEnv
 globalSetup opts = do
   opts <- checkConfig opts
   ctMap <- liftIO $ newIORef M.empty
-  return $ GlobalEnv opts ctMap
+  ilMap <- liftIO $ newIORef S.empty
+  return $ GlobalEnv opts ctMap ilMap
 
 initCompileEnv :: GlobalEnv -> TopLevelModuleName -> SpecialRules -> CompileEnv
 initCompileEnv genv tlm rewrites = CompileEnv

src/Agda2Hs/Compile/Function.hs

@@ -396,31 +396,11 @@ checkTransparentPragma def = compileFun False def >>= \case
       "A transparent function must have exactly one non-erased argument and return it unchanged."
 
 
--- | Ensure a definition can be defined as inline.
+-- | Mark a definition as one that should be inlined.
 checkInlinePragma :: Definition -> C ()
-checkInlinePragma def@Defn{defName = f} = do
-  let Function{funClauses = cs} = theDef def
-  case filter (isJust . clauseBody) cs of
-    [c] ->
-      unlessM (allowedPats (namedClausePats c)) $ agda2hsErrorM $
-        "Cannot make function" <+> prettyTCM (defName def) <+> "inlinable." <+>
-        "Inline functions can only use variable patterns or transparent record constructor patterns."
-    _ ->
-      agda2hsErrorM $
-        "Cannot make function" <+> prettyTCM f <+> "inlinable." <+>
-        "An inline function must have exactly one clause."
-
-  where allowedPat :: DeBruijnPattern -> C Bool
-        allowedPat VarP{} = pure True
-        -- only allow matching on (unboxed) record constructors
-        allowedPat (ConP ch ci cargs) =
-          isUnboxConstructor (conName ch) >>= \case
-            Just _  -> allowedPats cargs
-            Nothing -> pure False
-        allowedPat _ = pure False
-
-        allowedPats :: NAPs -> C Bool
-        allowedPats pats = allM (allowedPat . dget . dget) pats
+checkInlinePragma def@(Defn { defName = q , theDef = df }) = do
+  let qs = fromMaybe [] $ getMutual_ df
+  addInlineSymbols $ q : qs
 
 checkCompileToFunctionPragma :: Definition -> String -> C ()
 checkCompileToFunctionPragma def s = noCheckNames $ do

src/Agda2Hs/Compile/Term.hs

@@ -171,8 +171,8 @@ compileDef f ty args | Just sem <- isSpecialDef f = do
   sem ty args
 
 compileDef f ty args =
-  ifM (isTransparentFunction f) (compileErasedApp ty args) $
-  ifM (isInlinedFunction f) (compileInlineFunctionApp f ty args) $ do
+  ifM (isTransparentFunction f) (compileErasedApp ty args) $ do
+
     reportSDoc "agda2hs.compile.term" 12 $ text "Compiling application of regular function:" <+> prettyTCM f
 
     let defMod = qnameModule f
@@ -458,14 +458,18 @@ compileTerm ty v = do
 
   v <- instantiate v
 
+  toInline <- getInlineSymbols
+  v <- locallyReduceDefs (OnlyReduceDefs toInline) $ reduce v
+
   let bad s t = agda2hsErrorM $ vcat
         [ text "cannot compile" <+> text (s ++ ":")
         , nest 2 $ prettyTCM t
         ]
 
   reduceProjectionLike v >>= \case
 
-    Def f es -> do
+    v@(Def f es) -> do
+      whenM (isInlinedFunction f) $ bad "inlined function" v
       ty <- defType <$> getConstInfo f
       compileSpined (compileDef f ty) (Def f) ty es
 

src/Agda2Hs/Compile/Types.hs

@@ -29,6 +29,8 @@ data GlobalEnv = GlobalEnv
   { globalOptions :: Options
   , compileToMap :: IORef (Map QName QName)
   -- ^ names with a compile-to pragma
+  , inlineSymbols :: IORef (Set QName)
+  -- ^ names of functions that should be inlined
   }
 
 type ModuleEnv   = TopLevelModuleName

src/Agda2Hs/Compile/Utils.hs

@@ -12,6 +12,8 @@ import Data.List ( isPrefixOf, stripPrefix )
 import Data.Maybe ( isJust )
 import qualified Data.Map as M
 import Data.String ( IsString(..) )
+import Data.Set ( Set )
+import qualified Data.Set as S
 
 import GHC.Stack (HasCallStack)
 
@@ -281,8 +283,28 @@ isTupleProjection q =
 isTransparentFunction :: QName -> C Bool
 isTransparentFunction q = (== TransparentPragma) <$> getPragma q
 
+getInlineSymbols :: C (Set QName)
+getInlineSymbols = do
+  ilSetRef <- asks $ inlineSymbols . globalEnv
+  liftIO $ readIORef ilSetRef
+
+debugInlineSymbols :: C ()
+debugInlineSymbols = do
+  ilSetRef <- asks $ inlineSymbols . globalEnv
+  ilSet <- liftIO $ readIORef ilSetRef
+  reportSDoc "agda2hs.compile.inline" 50 $ text $
+    show $ map prettyShow $ S.toList ilSet
+
 isInlinedFunction :: QName -> C Bool
-isInlinedFunction q = (== InlinePragma) <$> getPragma q
+isInlinedFunction q = S.member q <$> getInlineSymbols
+
+addInlineSymbols :: [QName] -> C ()
+addInlineSymbols qs = do
+  reportSDoc "agda2hs.compile.inline" 15 $
+    "Adding inline rules for" <+> pretty qs
+  ilSetRef <- asks $ inlineSymbols . globalEnv
+  liftIO $ modifyIORef ilSetRef $ \s -> foldr S.insert s qs
+
 
 debugCompileToMap :: C ()
 debugCompileToMap = do

test/Fail/Inline.agda

@@ -6,3 +6,8 @@ tail' : List a → List a
 tail' (x ∷ xs) = xs
 tail' [] = []
 {-# COMPILE AGDA2HS tail' inline #-}
+
+test : List a → List a
+test = tail'
+
+{-# COMPILE AGDA2HS test #-}

test/Fail/Inline2.agda

@@ -5,3 +5,8 @@ open import Haskell.Prelude
 tail' : (xs : List a) → @0 {{ NonEmpty xs }}  → List a
 tail' (x ∷ xs) = xs
 {-# COMPILE AGDA2HS tail' inline #-}
+
+test : (xs : List a) → @0 {{ NonEmpty xs }}  → List a
+test = tail'
+
+{-# COMPILE AGDA2HS test #-}

test/golden/Inline.err

@@ -1,3 +1,4 @@
-test/Fail/Inline.agda:5.1-6: error: [CustomBackendError]
+test/Fail/Inline.agda:10.1-5: error: [CustomBackendError]
 agda2hs:
-  Cannot make function tail' inlinable. An inline function must have exactly one clause.
+  cannot compile inlined function:
+    tail'

test/golden/Inline2.err

@@ -1,3 +1,4 @@
-test/Fail/Inline2.agda:5.1-6: error: [CustomBackendError]
+test/Fail/Inline2.agda:9.1-5: error: [CustomBackendError]
 agda2hs:
-  Cannot make function tail' inlinable. Inline functions can only use variable patterns or transparent record constructor patterns.
+  cannot compile inlined function:
+    tail'