Merge pull request #1336 from cypheon/lambdalift-opt

Lambdalift optimisation: drop unused variables from outer scopes

src/Compiler/LambdaLift.idr

@@ -158,6 +158,61 @@ unload fc _ f [] = pure f
 -- only outermost LApp must be lazy as rest will be closures
 unload fc lazy f (a :: as) = unload fc Nothing (LApp fc lazy f a) as
 
+record Used (vars : List Name) where
+  constructor MkUsed
+  used : Vect (length vars) Bool
+
+initUsed : {vars : _} -> Used vars
+initUsed {vars} = MkUsed (replicate (length vars) False)
+
+lengthDistributesOverAppend
+  : (xs, ys : List a)
+  -> length (xs ++ ys) = length xs + length ys
+lengthDistributesOverAppend [] ys = Refl
+lengthDistributesOverAppend (x :: xs) ys =
+  cong S $ lengthDistributesOverAppend xs ys
+
+weakenUsed : {outer : _} -> Used vars -> Used (outer ++ vars)
+weakenUsed {outer} (MkUsed xs) =
+  MkUsed (rewrite lengthDistributesOverAppend outer vars in
+         (replicate (length outer) False ++ xs))
+
+contractUsed : (Used (x::vars)) -> Used vars
+contractUsed (MkUsed xs) = MkUsed (tail xs)
+
+contractUsedMany : {remove : _} ->
+                   (Used (remove ++ vars)) ->
+                   Used vars
+contractUsedMany {remove=[]} x = x
+contractUsedMany {remove=(r::rs)} x = contractUsedMany {remove=rs} (contractUsed x)
+
+markUsed : {vars : _} ->
+           (idx : Nat) ->
+           {0 prf : IsVar x idx vars} ->
+           Used vars ->
+           Used vars
+markUsed {vars} {prf} idx (MkUsed us) =
+  let newUsed = replaceAt (finIdx prf) True us in
+  MkUsed newUsed
+    where
+    finIdx : {vars : _} -> {idx : _} ->
+               (0 prf : IsVar x idx vars) ->
+               Fin (length vars)
+    finIdx {idx=Z} First = FZ
+    finIdx {idx=S x} (Later l) = FS (finIdx l)
+
+getUnused : Used vars ->
+            Vect (length vars) Bool
+getUnused (MkUsed uv) = map not uv
+
+total
+dropped : (vars : List Name) ->
+          (drop : Vect (length vars) Bool) ->
+          List Name
+dropped [] _ = []
+dropped (x::xs) (False::us) = x::(dropped xs us)
+dropped (x::xs) (True::us) = dropped xs us
+
 mutual
   makeLam : {auto l : Ref Lifts LDefs} ->
             {vars : _} ->
@@ -168,24 +223,27 @@ mutual
   makeLam fc bound (CLam _ x sc') = makeLam fc {doLazyAnnots} {lazy} (x :: bound) sc'
   makeLam {vars} fc bound sc
       = do scl <- liftExp {doLazyAnnots} {lazy} sc
+           -- Find out which variables aren't used in the new definition, and
+           -- do not abstract over them in the new definition.
+           let scUsedL = usedVars initUsed scl
+               unusedContracted = contractUsedMany {remove=bound} scUsedL
+               unused = getUnused unusedContracted
+               scl' = dropUnused {outer=bound} unused scl
            n <- genName
            ldefs <- get Lifts
-           put Lifts (record { defs $= ((n, MkLFun vars bound scl) ::) } ldefs)
-           -- TODO: an optimisation here would be to spot which variables
-           -- aren't used in the new definition, and not abstract over them
-           -- in the new definition. Given that we have to do some messing
-           -- about with indices anyway, it's probably not costly to do.
-           pure $ LUnderApp fc n (length bound) (allVars fc vars)
+           put Lifts (record { defs $= ((n, MkLFun (dropped vars unused) bound scl') ::) } ldefs)
+           pure $ LUnderApp fc n (length bound) (allVars fc vars unused)
     where
-        allPrfs : (vs : List Name) -> List (Var vs)
-        allPrfs [] = []
-        allPrfs (v :: vs) = MkVar First :: map weaken (allPrfs vs)
+        allPrfs : (vs : List Name) -> (unused : Vect (length vs) Bool) -> List (Var vs)
+        allPrfs [] _ = []
+        allPrfs (v :: vs) (False::uvs) = MkVar First :: map weaken (allPrfs vs uvs)
+        allPrfs (v :: vs) (True::uvs) = map weaken (allPrfs vs uvs)
 
         -- apply to all the variables. 'First' will be first in the last, which
         -- is good, because the most recently bound name is the first argument to
         -- the resulting function
-        allVars : FC -> (vs : List Name) -> List (Lifted vs)
-        allVars fc vs = map (\ (MkVar p) => LLocal fc p) (allPrfs vs)
+        allVars : FC -> (vs : List Name) -> (unused : Vect (length vs) Bool) -> List (Lifted vs)
+        allVars fc vs unused = map (\ (MkVar p) => LLocal fc p) (allPrfs vs unused)
 
 -- if doLazyAnnots = True then annotate function application with laziness
 -- otherwise use old behaviour (thunk is a function)
@@ -234,6 +292,117 @@ mutual
   liftExp (CErased fc) = pure $ LErased fc
   liftExp (CCrash fc str) = pure $ LCrash fc str
 
+  usedVars : {vars : _} ->
+             {auto l : Ref Lifts LDefs} ->
+             Used vars ->
+             Lifted vars ->
+             Used vars
+  usedVars used (LLocal {idx} fc prf) =
+    markUsed {prf} idx used
+  usedVars used (LAppName fc lazy n args) =
+    foldl (usedVars {vars}) used args
+  usedVars used (LUnderApp fc n miss args) =
+    foldl (usedVars {vars}) used args
+  usedVars used (LApp fc lazy c arg) =
+    usedVars (usedVars used arg) c
+  usedVars used (LLet fc x val sc) =
+    let innerUsed = contractUsed $ usedVars (weakenUsed {outer=[x]} used) sc in
+        usedVars innerUsed val
+  usedVars used (LCon fc n tag args) =
+    foldl (usedVars {vars}) used args
+  usedVars used (LOp fc lazy fn args) =
+    foldl (usedVars {vars}) used args
+  usedVars used (LExtPrim fc lazy fn args) =
+    foldl (usedVars {vars}) used args
+  usedVars used (LConCase fc sc alts def) =
+      let defUsed = maybe used (usedVars used {vars}) def
+          scDefUsed = usedVars defUsed sc in
+          foldl usedConAlt scDefUsed alts
+    where
+      usedConAlt : {default Nothing lazy : Maybe LazyReason} ->
+                   Used vars -> LiftedConAlt vars -> Used vars
+      usedConAlt used (MkLConAlt n tag args sc) =
+        contractUsedMany {remove=args} (usedVars (weakenUsed used) sc)
+
+  usedVars used (LConstCase fc sc alts def) =
+      let defUsed = maybe used (usedVars used {vars}) def
+          scDefUsed = usedVars defUsed sc in
+          foldl usedConstAlt scDefUsed alts
+    where
+      usedConstAlt : {default Nothing lazy : Maybe LazyReason} ->
+                     Used vars -> LiftedConstAlt vars -> Used vars
+      usedConstAlt used (MkLConstAlt c sc) = usedVars used sc
+
+  usedVars used (LPrimVal _ _) = used
+  usedVars used (LErased _) = used
+  usedVars used (LCrash _ _) = used
+
+  dropIdx : {vars : _} ->
+            {idx : _} ->
+            (outer : List Name) ->
+            (unused : Vect (length vars) Bool) ->
+            (0 p : IsVar x idx (outer ++ vars)) ->
+            Var (outer ++ (dropped vars unused))
+  dropIdx [] (False::_) First = MkVar First
+  dropIdx [] (True::_) First = assert_total $
+    idris_crash "INTERNAL ERROR: Referenced variable marked as unused"
+  dropIdx [] (False::rest) (Later p) = Var.later $ dropIdx [] rest p
+  dropIdx [] (True::rest) (Later p) = dropIdx [] rest p
+  dropIdx (_::xs) unused First = MkVar First
+  dropIdx (_::xs) unused (Later p) = Var.later $ dropIdx xs unused p
+
+  dropUnused : {vars : _} ->
+               {auto l : Ref Lifts LDefs} ->
+               {outer : List Name} ->
+               (unused : Vect (length vars) Bool) ->
+               (l : Lifted (outer ++ vars)) ->
+               Lifted (outer ++ (dropped vars unused))
+  dropUnused _ (LPrimVal fc val) = LPrimVal fc val
+  dropUnused _ (LErased fc) = LErased fc
+  dropUnused _ (LCrash fc msg) = LCrash fc msg
+  dropUnused {outer} unused (LLocal fc p) =
+    let (MkVar p') = dropIdx outer unused p in LLocal fc p'
+  dropUnused unused (LCon fc n tag args) =
+    let args' = map (dropUnused unused) args in
+        LCon fc n tag args'
+  dropUnused {outer} unused (LLet fc n val sc) =
+    let val' = dropUnused unused val
+        sc' = dropUnused {outer=n::outer} (unused) sc in
+        LLet fc n val' sc'
+  dropUnused unused (LApp fc lazy c arg) =
+    let c' = dropUnused unused c
+        arg' = dropUnused unused arg in
+        LApp fc lazy c' arg'
+  dropUnused unused (LOp fc lazy fn args) =
+    let args' = map (dropUnused unused) args in
+        LOp fc lazy fn args'
+  dropUnused unused (LExtPrim fc lazy n args) =
+    let args' = map (dropUnused unused) args in
+        LExtPrim fc lazy n args'
+  dropUnused unused (LAppName fc lazy n args) =
+    let args' = map (dropUnused unused) args in
+        LAppName fc lazy n args'
+  dropUnused unused (LUnderApp fc n miss args) =
+    let args' = map (dropUnused unused) args in
+        LUnderApp fc n miss args'
+  dropUnused {vars} {outer} unused (LConCase fc sc alts def) =
+    let alts' = map dropConCase alts in
+        LConCase fc (dropUnused unused sc) alts' (map (dropUnused unused) def)
+    where
+      dropConCase : LiftedConAlt (outer ++ vars) ->
+                    LiftedConAlt (outer ++ (dropped vars unused))
+      dropConCase (MkLConAlt n t args sc) =
+        let sc' = (rewrite sym $ appendAssociative args outer vars in sc)
+            droppedSc = dropUnused {vars=vars} {outer=args++outer} unused sc' in
+        MkLConAlt n t args (rewrite appendAssociative args outer (dropped vars unused) in droppedSc)
+  dropUnused {vars} {outer} unused (LConstCase fc sc alts def) =
+    let alts' = map dropConstCase alts in
+        LConstCase fc (dropUnused unused sc) alts' (map (dropUnused unused) def)
+    where
+      dropConstCase : LiftedConstAlt (outer ++ vars) ->
+                      LiftedConstAlt (outer ++ (dropped vars unused))
+      dropConstCase (MkLConstAlt c val) = MkLConstAlt c (dropUnused unused val)
+
 export
 liftBody : {vars : _} -> {doLazyAnnots : Bool} ->
            Name -> CExp vars -> Core (Lifted vars, List (Name, LiftedDef))

src/Core/Core.idr

@@ -566,6 +566,11 @@ Catchable Core Error where
                          Right val => pure (Right val))
   throw = coreFail
 
+-- Prelude.Monad.foldlM hand specialised for Core
+export
+foldlC : Foldable t => (a -> b -> Core a) -> a -> t b -> Core a
+foldlC fm a0 = foldl (\ma,b => ma >>= flip fm b) (pure a0)
+
 -- Traversable (specialised)
 traverse' : (a -> Core b) -> List a -> List b -> Core (List b)
 traverse' f [] acc = pure (reverse acc)

src/Idris/Package.idr

@@ -415,10 +415,6 @@ bitraverseC f g (This a)   = [| This (f a) |]
 bitraverseC f g (That b)   = [| That (g b) |]
 bitraverseC f g (Both a b) = [| Both (f a) (g b) |]
 
--- Prelude.Monad.foldlM hand specialised for Core
-foldlC : Foldable t => (a -> b -> Core a) -> a -> t b -> Core a
-foldlC fm a0 = foldl (\ma,b => ma >>= flip fm b) (pure a0)
-
 -- Data.StringTrie.foldWithKeysM hand specialised for Core
 foldWithKeysC : Monoid b => (List String -> Core b) -> (List String -> a -> Core b) -> StringTrie a -> Core b
 foldWithKeysC {a} {b} fk fv = go []