Calculate whether an inlining is safe

It's safe if all top level arguments are used at most once, meaning that
there's no risk of duplication.

Author: edwinb

libs/prelude/Prelude/Interfaces.idr

@@ -103,6 +103,7 @@ public export
 ($>) fa b = map (const b) fa
 
 ||| Run something for effects, throwing away the return value.
+%inline
 public export
 ignore : Functor f => f a -> f ()
 ignore = map (const ())

src/TTImp/Elab/Case.idr

@@ -1,5 +1,6 @@
 module TTImp.Elab.Case
 
+import Core.CaseTree
 import Core.Context
 import Core.Context.Log
 import Core.Core
@@ -226,13 +227,6 @@ caseBlock {vars} rigc elabinfo fc nest env scr scrtm scrty caseRig alts expected
          setFlag fc (Resolved cidx) (SetTotal PartialOK)
          let caseRef : Term vars = Ref fc Func (Resolved cidx)
 
-         -- If there's no duplication of the scrutinee in the block,
-         -- inline it.
-         -- This will be the case either if the scrutinee is a variable, in
-         -- which case the duplication won't hurt, or if (TODO) none of the
-         -- case patterns in alts are just a variable
-         maybe (pure ()) (const (setFlag fc casen Inline)) splitOn
-
          let applyEnv = applyToFull fc caseRef env
          let appTm : Term vars
                    = maybe (App fc applyEnv scrtm)
@@ -254,6 +248,17 @@ caseBlock {vars} rigc elabinfo fc nest env scr scrtm scrty caseRig alts expected
          let olddelayed = delayedElab ust
          put UST (record { delayedElab = [] } ust)
          processDecl [InCase] nest' [] (IDef fc casen alts')
+
+         -- If there's no duplication of the scrutinee in the block,
+         -- flag it as inlinable.
+         -- This will be the case either if the scrutinee is a variable, in
+         -- which case the duplication won't hurt, or if there's no variable
+         -- duplicated in the body (what ghc calls W-safe)
+         -- We'll check that second condition later, after generating the
+         -- runtime (erased) case trees
+         let inlineOK = maybe False (const True) splitOn
+         when inlineOK $ setFlag fc casen Inline
+
          ust <- get UST
          put UST (record { delayedElab = olddelayed } ust)
 

src/TTImp/Elab/Utils.idr

@@ -1,5 +1,6 @@
 module TTImp.Elab.Utils
 
+import Core.CaseTree
 import Core.Context
 import Core.Core
 import Core.Env
@@ -116,3 +117,169 @@ bindReq {vs = n :: _} fc (b :: env) (KeepCons p) ns tm
             (Bind fc _ (Pi (binderLoc b) (multiplicity b) Explicit (binderType b')) tm)
 bindReq fc (b :: env) (DropCons p) ns tm
     = bindReq fc env p ns tm
+
+-- This machinery is to calculate whether any top level argument is used
+-- more than once in a case block, in which case inlining wouldn't be safe
+-- since it might duplicate work.
+
+data ArgUsed = Used1 -- been used
+             | Used0 -- not used
+             | LocalVar -- don't care if it's used
+
+data Usage : List Name -> Type where
+     Nil : Usage []
+     (::) : ArgUsed -> Usage xs -> Usage (x :: xs)
+
+initUsed : (xs : List Name) -> Usage xs
+initUsed [] = []
+initUsed (x :: xs) = Used0 :: initUsed xs
+
+initUsedCase : (xs : List Name) -> Usage xs
+initUsedCase [] = []
+initUsedCase [x] = [Used0]
+initUsedCase (x :: xs) = LocalVar :: initUsedCase xs
+
+setUsedVar : {idx : _} ->
+             (0 _ : IsVar n idx xs) -> Usage xs -> Usage xs
+setUsedVar First (Used0 :: us) = Used1 :: us
+setUsedVar (Later p) (x :: us) = x :: setUsedVar p us
+setUsedVar First us = us
+
+isUsed : {idx : _} ->
+         (0 _ : IsVar n idx xs) -> Usage xs -> Bool
+isUsed First (Used1 :: us) = True
+isUsed First (_ :: us) = False
+isUsed (Later p) (_ :: us) = isUsed p us
+
+data Used : Type where
+
+setUsed : {idx : _} ->
+          {auto u : Ref Used (Usage vars)} ->
+          (0 _ : IsVar n idx vars) -> Core ()
+setUsed p
+    = do used <- get Used
+         put Used (setUsedVar p used)
+
+extendUsed : ArgUsed -> (new : List Name) -> Usage vars -> Usage (new ++ vars)
+extendUsed a [] x = x
+extendUsed a (y :: xs) x = a :: extendUsed a xs x
+
+dropUsed : (new : List Name) -> Usage (new ++ vars) -> Usage vars
+dropUsed [] x = x
+dropUsed (x :: xs) (u :: us) = dropUsed xs us
+
+inExtended : ArgUsed -> (new : List Name) ->
+             {auto u : Ref Used (Usage vars)} ->
+             (Ref Used (Usage (new ++ vars)) -> Core a) ->
+             Core a
+inExtended a new sc
+    = do used <- get Used
+         u' <- newRef Used (extendUsed a new used)
+         res <- sc u'
+         put Used (dropUsed new !(get Used @{u'}))
+         pure res
+
+termInlineSafe : {vars : _} ->
+                 {auto u : Ref Used (Usage vars)} ->
+                 Term vars -> Core Bool
+termInlineSafe (Local fc isLet idx p)
+   = if isUsed p !(get Used)
+        then pure False
+         else do setUsed p
+                 pure True
+termInlineSafe (Meta fc x y xs)
+    = allInlineSafe xs
+  where
+    allInlineSafe : List (Term vars) -> Core Bool
+    allInlineSafe [] = pure True
+    allInlineSafe (x :: xs)
+        = do xok <- termInlineSafe x
+             if xok
+                then allInlineSafe xs
+                else pure False
+termInlineSafe (Bind fc x b scope)
+   = do bok <- binderInlineSafe b
+        if bok
+           then inExtended LocalVar [x] (\u' => termInlineSafe scope)
+           else pure False
+  where
+    binderInlineSafe : Binder (Term vars) -> Core Bool
+    binderInlineSafe (Let _ _ val _) = termInlineSafe val
+    binderInlineSafe _ = pure True
+termInlineSafe (App fc fn arg)
+    = do fok <- termInlineSafe fn
+         if fok
+            then termInlineSafe arg
+            else pure False
+termInlineSafe (As fc x as pat) = termInlineSafe pat
+termInlineSafe (TDelayed fc x ty) = termInlineSafe ty
+termInlineSafe (TDelay fc x ty arg) = termInlineSafe arg
+termInlineSafe (TForce fc x val) = termInlineSafe val
+termInlineSafe _ = pure True
+
+mutual
+  caseInlineSafe : {vars : _} ->
+                   {auto u : Ref Used (Usage vars)} ->
+                   CaseTree vars -> Core Bool
+  caseInlineSafe (Case idx p scTy xs)
+      = if isUsed p !(get Used)
+           then pure False
+           else do setUsed p
+                   altsSafe xs
+    where
+      altsSafe : List (CaseAlt vars) -> Core Bool
+      altsSafe [] = pure True
+      altsSafe (a :: as)
+          = do u <- get Used
+               aok <- caseAltInlineSafe a
+               if aok
+                  then do -- We can reset the usage information, because we're
+                          -- only going to use one alternative at a time
+                          put Used u
+                          altsSafe as
+                  else pure False
+  caseInlineSafe (STerm x tm) = termInlineSafe tm
+  caseInlineSafe (Unmatched msg) = pure True
+  caseInlineSafe Impossible = pure True
+
+  caseAltInlineSafe : {vars : _} ->
+                      {auto u : Ref Used (Usage vars)} ->
+                      CaseAlt vars -> Core Bool
+  caseAltInlineSafe (ConCase x tag args sc)
+      = inExtended Used0 args (\u' => caseInlineSafe sc)
+  caseAltInlineSafe (DelayCase ty arg sc)
+      = inExtended Used0 [ty, arg] (\u' => caseInlineSafe sc)
+  caseAltInlineSafe (ConstCase x sc) = caseInlineSafe sc
+  caseAltInlineSafe (DefaultCase sc) = caseInlineSafe sc
+
+-- An inlining is safe if no variable is used more than once in the tree,
+-- which means that there's no risk of an input being evaluated more than
+-- once after the definition is expanded.
+export
+inlineSafe : {vars : _} ->
+             CaseTree vars -> Core Bool
+inlineSafe t
+    = do u <- newRef Used (initUsed vars)
+         caseInlineSafe t
+
+export
+canInlineDef : {auto c : Ref Ctxt Defs} ->
+               Name -> Core Bool
+canInlineDef n
+    = do defs <- get Ctxt
+         Just (PMDef _ _ _ rtree _) <- lookupDefExact n (gamma defs)
+             | _ => pure False
+         inlineSafe rtree
+
+-- This is a special case because the only argument we actually care about
+-- is the last one, since the others are just variables passed through from
+-- the environment, and duplicating a variable doesn't cost anything.
+export
+canInlineCaseBlock : {auto c : Ref Ctxt Defs} ->
+                     Name -> Core Bool
+canInlineCaseBlock n
+    = do defs <- get Ctxt
+         Just (PMDef _ vars _ rtree _) <- lookupDefExact n (gamma defs)
+             | _ => pure False
+         u <- newRef Used (initUsedCase vars)
+         caseInlineSafe rtree

src/TTImp/ProcessDef.idr

@@ -750,7 +750,22 @@ mkRunTime fc n
            ignore $ addDef n $
                        record { definition = PMDef r rargs tree_ct tree_rt pats
                               } gdef
+           -- If it's a case block, and not already set as inlinable,
+           -- check if it's safe to inline
+           when (caseName !(toFullNames n) && noInline (flags gdef)) $
+             do inl <- canInlineCaseBlock n
+                when inl $ setFlag fc n Inline
   where
+    noInline : List DefFlag -> Bool
+    noInline (Inline :: _) = False
+    noInline (x :: xs) = noInline xs
+    noInline _ = True
+
+    caseName : Name -> Bool
+    caseName (CaseBlock _ _) = True
+    caseName (NS _ n) = caseName n
+    caseName _ = False
+
     mkCrash : {vars : _} -> String -> Term vars
     mkCrash msg
        = apply fc (Ref fc Func (NS builtinNS (UN "idris_crash")))

tests/codegen/builtin001/expected

@@ -5,10 +5,8 @@ prim__mul_Integer = [{arg:N}, {arg:N}]: (*Integer [!{arg:N}, !{arg:N}])
 Main.plus = [{arg:N}, {arg:N}]: (%case !{arg:N} [(%constcase 0 !{arg:N})] Just (%let {e:N} (-Integer [!{arg:N}, 1]) (+Integer [1, (Main.plus [!{e:N}, !{arg:N}])])))
 Main.main = [{ext:N}]: (Main.plus [(+Integer [1, 0]), (+Integer [1, (+Integer [1, 0])])])
 Builtin.believe_me = [{ext:N}]: (believe_me [___, ___, !{ext:N}])
-Prelude.Types.case block in prim__integerToNat = [{arg:N}, {arg:N}]: (%case !{arg:N} [(%constcase 1 (Builtin.believe_me [!{arg:N}])), (%constcase 0 0)] Nothing)
-Prelude.Types.prim__integerToNat = [{arg:N}]: (Prelude.Types.case block in prim__integerToNat [!{arg:N}, (%case (<=Integer [0, !{arg:N}]) [(%constcase 0 0)] Just 1)])
-PrimIO.case block in unsafePerformIO = [{arg:N}, {arg:N}]: (PrimIO.unsafeDestroyWorld [___, !{arg:N}])
-PrimIO.unsafePerformIO = [{arg:N}]: (PrimIO.unsafeCreateWorld [(%lam w (PrimIO.case block in unsafePerformIO [!{arg:N}, (!{arg:N} [!w])]))])
+Prelude.Types.prim__integerToNat = [{arg:N}]: (%case (%case (<=Integer [0, !{arg:N}]) [(%constcase 0 0)] Just 1) [(%constcase 1 (Builtin.believe_me [!{arg:N}])), (%constcase 0 0)] Nothing)
+PrimIO.unsafePerformIO = [{arg:N}]: (PrimIO.unsafeCreateWorld [(%lam w (PrimIO.unsafeDestroyWorld [___, (!{arg:N} [!w])]))])
 PrimIO.unsafeDestroyWorld = [{arg:N}, {arg:N}]: !{arg:N}
 PrimIO.unsafeCreateWorld = [{arg:N}]: (!{arg:N} [%MkWorld])