[ refactor ] ScopedSnocList: Swap Scope on SnocList (Phase 2)

idris2api.ipkg

@@ -198,6 +198,7 @@ modules =
     Libraries.Data.SnocList.HasLength,
     Libraries.Data.SnocList.LengthMatch,
     Libraries.Data.SnocList.SizeOf,
+    Libraries.Data.SnocList.Quantifiers.Extra,
     Libraries.Data.Span,
     Libraries.Data.SparseMatrix,
     Libraries.Data.String,

src/Compiler/ANF.idr

@@ -5,9 +5,13 @@ import Compiler.LambdaLift
 import Core.CompileExpr
 import Core.Context
 
+import Data.SnocList.Quantifiers
+import Data.String
 import Data.SortedSet
 import Data.Vect
 
+import Libraries.Data.SnocList.Extra
+
 %default covering
 
 -- Convert the lambda lifted form to ANF, with variable names made explicit.
@@ -82,26 +86,26 @@ mutual
   Show ANF where
     show (AV _ v) = show v
     show (AAppName fc lazy n args)
-        = show n ++ showLazy lazy ++ "(" ++ showSep ", " (map show args) ++ ")"
+        = show n ++ showLazy lazy ++ "(" ++ joinBy ", " (map show args) ++ ")"
     show (AUnderApp fc n m args)
         = "<" ++ show n ++ " underapp " ++ show m ++ ">(" ++
-          showSep ", " (map show args) ++ ")"
+          joinBy ", " (map show args) ++ ")"
     show (AApp fc lazy c arg)
         = show c ++ showLazy lazy ++ " @ (" ++ show arg ++ ")"
     show (ALet fc x val sc)
         = "%let v" ++ show x ++ " = (" ++ show val ++ ") in (" ++ show sc ++ ")"
     show (ACon fc n _ t args)
-        = "%con " ++ show n ++ "(" ++ showSep ", " (map show args) ++ ")"
+        = "%con " ++ show n ++ "(" ++ joinBy ", " (map show args) ++ ")"
     show (AOp fc lazy op args)
-        = "%op " ++ show op ++ showLazy lazy ++ "(" ++ showSep ", " (toList (map show args)) ++ ")"
+        = "%op " ++ show op ++ showLazy lazy ++ "(" ++ joinBy ", " (toList (map show args)) ++ ")"
     show (AExtPrim fc lazy p args)
-        = "%extprim " ++ show p ++ showLazy lazy ++ "(" ++ showSep ", " (map show args) ++ ")"
+        = "%extprim " ++ show p ++ showLazy lazy ++ "(" ++ joinBy ", " (map show args) ++ ")"
     show (AConCase fc sc alts def)
         = "%case " ++ show sc ++ " of { "
-             ++ showSep "| " (map show alts) ++ " " ++ show def ++ " }"
+             ++ joinBy "| " (map show alts) ++ " " ++ show def ++ " }"
     show (AConstCase fc sc alts def)
         = "%case " ++ show sc ++ " of { "
-             ++ showSep "| " (map show alts) ++ " " ++ show def ++ " }"
+             ++ joinBy "| " (map show alts) ++ " " ++ show def ++ " }"
     show (APrimVal _ x) = show x
     show (AErased _) = "___"
     show (ACrash _ x) = "%CRASH(" ++ show x ++ ")"
@@ -111,7 +115,7 @@ mutual
   Show AConAlt where
     show (MkAConAlt n _ t args sc)
         = "%conalt " ++ show n ++
-             "(" ++ showSep ", " (map showArg args) ++ ") => " ++ show sc
+             "(" ++ joinBy ", " (map showArg args) ++ ") => " ++ show sc
       where
         showArg : Int -> String
         showArg i = "v" ++ show i
@@ -136,6 +140,12 @@ Show ANFDef where
 AVars : Scope -> Type
 AVars = All (\_ => Int)
 
+namespace AVars
+  public export
+  empty : AVars Scope.empty
+  empty = [<]
+
+
 data Next : Type where
 
 nextVar : {auto v : Ref Next Int} ->
@@ -145,10 +155,6 @@ nextVar
          put Next (i + 1)
          pure i
 
-lookup : {idx : _} -> (0 p : IsVar x idx vs) -> AVars vs -> Int
-lookup First (x :: xs) = x
-lookup (Later p) (x :: xs) = lookup p xs
-
 bindArgs : {auto v : Ref Next Int} ->
            List ANF -> Core (List (AVar, Maybe ANF))
 bindArgs [] = pure []
@@ -183,6 +189,15 @@ mlet fc val sc
     = do i <- nextVar
          pure $ ALet fc i val (sc (ALocal i))
 
+bindAsFresh :
+  {auto v : Ref Next Int} ->
+  (args : List Name) -> AVars vars' ->
+  Core (List Int, AVars (Scope.ext vars' args))
+bindAsFresh [] vs = pure ([], vs)
+bindAsFresh (n :: ns) vs
+    = do i <- nextVar
+         mapFst (i ::) <$> bindAsFresh ns (vs :< i)
+
 mutual
   anfArgs : {auto v : Ref Next Int} ->
             FC -> AVars vars ->
@@ -193,7 +208,7 @@ mutual
 
   anf : {auto v : Ref Next Int} ->
         AVars vars -> Lifted vars -> Core ANF
-  anf vs (LLocal fc p) = pure $ AV fc (ALocal (lookup p vs))
+  anf vs (LLocal fc p) = pure $ AV fc (ALocal (lookup vs p))
   anf vs (LAppName fc lazy n args)
       = anfArgs fc vs args (AAppName fc lazy n)
   anf vs (LUnderApp fc n m args)
@@ -205,7 +220,7 @@ mutual
                   _ => ACrash fc "Can't happen (AApp)"
   anf vs (LLet fc x val sc)
       = do i <- nextVar
-           let vs' = i :: vs
+           let vs' = vs :< i
            pure $ ALet fc i !(anf vs val) !(anf vs' sc)
   anf vs (LCon fc n ci t args)
       = anfArgs fc vs args (ACon fc n ci t)
@@ -234,16 +249,8 @@ mutual
   anfConAlt : {auto v : Ref Next Int} ->
               AVars vars -> LiftedConAlt vars -> Core AConAlt
   anfConAlt vs (MkLConAlt n ci t args sc)
-      = do (is, vs') <- bindArgs args vs
+      = do (is, vs') <- bindAsFresh args vs
            pure $ MkAConAlt n ci t is !(anf vs' sc)
-    where
-      bindArgs : (args : List Name) -> AVars vars' ->
-                 Core (List Int, AVars (args ++ vars'))
-      bindArgs [] vs = pure ([], vs)
-      bindArgs (n :: ns) vs
-          = do i <- nextVar
-               (is, vs') <- bindArgs ns vs
-               pure (i :: is, i :: vs')
 
   anfConstAlt : {auto v : Ref Next Int} ->
                 AVars vars -> LiftedConstAlt vars -> Core AConstAlt
@@ -254,25 +261,18 @@ export
 toANF : LiftedDef -> Core ANFDef
 toANF (MkLFun args scope sc)
     = do v <- newRef Next (the Int 0)
-         (iargs, vsNil) <- bindArgs args []
-         let vs : AVars args = rewrite sym (appendNilRightNeutral args) in
-                                      vsNil
-         (iargs', vs) <- bindArgs scope vs
-         pure $ MkAFun (iargs ++ reverse iargs') !(anf vs sc)
-  where
-    bindArgs : {auto v : Ref Next Int} ->
-               (args : List Name) -> AVars vars' ->
-               Core (List Int, AVars (args ++ vars'))
-    bindArgs [] vs = pure ([], vs)
-    bindArgs (n :: ns) vs
-        = do i <- nextVar
-             (is, vs') <- bindArgs ns vs
-             pure (i :: is, i :: vs')
+         (iargs, vsNil) <- bindAsFresh args AVars.empty
+         (iargs', vs) <- bindAsFresh (toList scope) vsNil
+         sc' <- anf vs $
+            do rewrite fishAsSnocAppend (cast args) (toList scope)
+               rewrite castToList scope
+               sc
+         pure $ MkAFun (iargs ++ iargs') sc'
 toANF (MkLCon t a ns) = pure $ MkACon t a ns
 toANF (MkLForeign ccs fargs t) = pure $ MkAForeign ccs fargs t
 toANF (MkLError err)
     = do v <- newRef Next (the Int 0)
-         pure $ MkAError !(anf [] err)
+         pure $ MkAError !(anf AVars.empty err)
 
 export
 freeVariables : ANF -> SortedSet AVar

src/Compiler/CaseOpts.idr

@@ -8,6 +8,8 @@ import Core.Context
 import Data.Vect
 
 import Libraries.Data.List.SizeOf
+import Libraries.Data.SnocList.SizeOf
+import Libraries.Data.SnocList.Extra
 
 %default covering
 
@@ -29,38 +31,38 @@ case t of
 
 shiftUnder : {args : _} ->
              {idx : _} ->
-             (0 p : IsVar n idx (x :: args ++ vars)) ->
-             NVar n (args ++ x :: vars)
+             (0 p : IsVar n idx (Scope.addInner vars (Scope.bind args x))) ->
+             NVar n (Scope.addInner (Scope.bind vars x) args)
 shiftUnder First = weakenNVar (mkSizeOf args) (MkNVar First)
 shiftUnder (Later p) = insertNVar (mkSizeOf args) (MkNVar p)
 
-shiftVar : {outer : Scope} -> {args : List Name} ->
-           NVar n (outer ++ (x :: args ++ vars)) ->
-           NVar n (outer ++ (args ++ x :: vars))
+shiftVar : {inner : Scope} -> {args : List Name} ->
+           NVar n ((vars <>< args :< x) ++ inner) ->
+           NVar n ((vars :< x <>< args) ++ inner)
 shiftVar nvar
-  = let out = mkSizeOf outer in
-    case locateNVar out nvar of
-      Left nvar => embed nvar
-      Right (MkNVar p) => weakenNs out (shiftUnder p)
+  = let inn = mkSizeOf inner in
+    case locateNVar inn nvar of
+      Left (MkNVar p) => weakenNs inn (shiftUndersN (mkSizeOf _) p)
+      Right nvar => embed nvar
 
 mutual
-  shiftBinder : {outer, args : _} ->
+  shiftBinder : {inner, args : _} ->
                 (new : Name) ->
-                CExp (outer ++ old :: (args ++ vars)) ->
-                CExp (outer ++ (args ++ new :: vars))
+                CExp (((vars <>< args) :< old) ++ inner) ->
+                CExp ((vars :< new <>< args) ++ inner)
   shiftBinder new (CLocal fc p)
       = case shiftVar (MkNVar p) of
              MkNVar p' => CLocal fc (renameVar p')
     where
-      renameVar : IsVar x i (outer ++ (args ++ (old :: rest))) ->
-                  IsVar x i (outer ++ (args ++ (new :: rest)))
+      renameVar : IsVar x i ((vars :< old <>< args) ++ local) ->
+                  IsVar x i ((vars :< new <>< args) ++ local)
       renameVar = believe_me -- it's the same index, so just the identity at run time
   shiftBinder new (CRef fc n) = CRef fc n
-  shiftBinder {outer} new (CLam fc n sc)
-      = CLam fc n $ shiftBinder {outer = n :: outer} new sc
+  shiftBinder {inner} new (CLam fc n sc)
+      = CLam fc n $ shiftBinder {inner = inner :< n} new sc
   shiftBinder new (CLet fc n inlineOK val sc)
       = CLet fc n inlineOK (shiftBinder new val)
-                           $ shiftBinder {outer = n :: outer} new sc
+                           $ shiftBinder {inner = inner :< n} new sc
   shiftBinder new (CApp fc f args)
       = CApp fc (shiftBinder new f) $ map (shiftBinder new) args
   shiftBinder new (CCon fc ci c tag args)
@@ -82,36 +84,36 @@ mutual
   shiftBinder new (CErased fc) = CErased fc
   shiftBinder new (CCrash fc msg) = CCrash fc msg
 
-  shiftBinderConAlt : {outer, args : _} ->
+  shiftBinderConAlt : {inner, args : _} ->
                 (new : Name) ->
-                CConAlt (outer ++ (x :: args ++ vars)) ->
-                CConAlt (outer ++ (args ++ new :: vars))
+                CConAlt (((vars <>< args) :< old) ++ inner) ->
+                CConAlt ((vars :< new <>< args) ++ inner)
   shiftBinderConAlt new (MkConAlt n ci t args' sc)
-      = let sc' : CExp ((args' ++ outer) ++ (x :: args ++ vars))
-                = rewrite sym (appendAssociative args' outer (x :: args ++ vars)) in sc in
+      = let sc' : CExp (((vars <>< args) :< old) ++ (inner <>< args'))
+                = rewrite sym $ snocAppendFishAssociative (vars <>< args :< old) inner args' in sc in
         MkConAlt n ci t args' $
-           rewrite (appendAssociative args' outer (args ++ new :: vars))
-             in shiftBinder new {outer = args' ++ outer} sc'
+               rewrite snocAppendFishAssociative (vars :< new <>< args) inner args' in
+               shiftBinder new sc'
 
-  shiftBinderConstAlt : {outer, args : _} ->
+  shiftBinderConstAlt : {inner, args : _} ->
                 (new : Name) ->
-                CConstAlt (outer ++ (x :: args ++ vars)) ->
-                CConstAlt (outer ++ (args ++ new :: vars))
+                CConstAlt (((vars <>< args) :< old) ++ inner) ->
+                CConstAlt ((vars :< new <>< args) ++ inner)
   shiftBinderConstAlt new (MkConstAlt c sc) = MkConstAlt c $ shiftBinder new sc
 
 -- If there's a lambda inside a case, move the variable so that it's bound
 -- outside the case block so that we can bind it just once outside the block
 liftOutLambda : {args : _} ->
                 (new : Name) ->
-                CExp (old :: args ++ vars) ->
-                CExp (args ++ new :: vars)
-liftOutLambda = shiftBinder {outer = Scope.empty}
+                CExp (Scope.bind (Scope.ext vars args) old) ->
+                CExp (Scope.ext (Scope.bind vars new) args)
+liftOutLambda = shiftBinder {inner = Scope.empty}
 
 -- If all the alternatives start with a lambda, we can have a single lambda
 -- binding outside
 tryLiftOut : (new : Name) ->
              List (CConAlt vars) ->
-             Maybe (List (CConAlt (new :: vars)))
+             Maybe (List (CConAlt (Scope.bind vars new)))
 tryLiftOut new [] = Just []
 tryLiftOut new (MkConAlt n ci t args (CLam fc x sc) :: as)
     = do as' <- tryLiftOut new as
@@ -121,7 +123,7 @@ tryLiftOut _ _ = Nothing
 
 tryLiftOutConst : (new : Name) ->
                   List (CConstAlt vars) ->
-                  Maybe (List (CConstAlt (new :: vars)))
+                  Maybe (List (CConstAlt (Scope.bind vars new)))
 tryLiftOutConst new [] = Just []
 tryLiftOutConst new (MkConstAlt c (CLam fc x sc) :: as)
     = do as' <- tryLiftOutConst new as
@@ -131,7 +133,7 @@ tryLiftOutConst _ _ = Nothing
 
 tryLiftDef : (new : Name) ->
              Maybe (CExp vars) ->
-             Maybe (Maybe (CExp (new :: vars)))
+             Maybe (Maybe (CExp (Scope.bind vars new)))
 tryLiftDef new Nothing = Just Nothing
 tryLiftDef new (Just (CLam fc x sc))
    = let sc' = liftOutLambda {args = []} new sc in
@@ -310,8 +312,8 @@ doCaseOfCase fc x xalts xdef alts def
     updateAlt (MkConAlt n ci t args sc)
         = MkConAlt n ci t args $
               CConCase fc sc
-                       (map (weakenNs (mkSizeOf args)) alts)
-                       (map (weakenNs (mkSizeOf args)) def)
+                       (map (weakensN (mkSizeOf args)) alts)
+                       (map (weakensN (mkSizeOf args)) def)
 
     updateDef : CExp vars -> CExp vars
     updateDef sc = CConCase fc sc alts def

src/Compiler/Common.idr

@@ -348,7 +348,7 @@ getCompileDataWith exports doLazyAnnots phase_in tm_in
                               traverse (lambdaLift doLazyAnnots) cseDefs
                          else pure []
 
-         let lifted = (mainname, MkLFun Scope.empty Scope.empty liftedtm) ::
+         let lifted = (mainname, MkLFun [] Scope.empty liftedtm) ::
                       (ldefs ++ concat lifted_in)
 
          anf <- if phase >= ANF