[ re #1282 ] Improve unelaboration of case expressions

src/Core/Env.idr

@@ -17,6 +17,11 @@ export
 extend : (x : Name) -> Binder (tm vars) -> Env tm vars -> Env tm (x :: vars)
 extend x = (::) {x}
 
+export
+(++) : {ns : _} -> Env Term ns -> Env Term vars -> Env Term (ns ++ vars)
+(++) (b :: bs) e = extend _ (map embed b) (bs ++ e)
+(++) [] e = e
+
 export
 length : Env tm xs -> Nat
 length [] = 0

src/TTImp/Unelab.idr

@@ -46,7 +46,7 @@ mutual
   ||| Unelaborate a call to a case expression as an inline case.
   ||| This should allow us to eventurally resugar case blocks and if-then-else calls.
   |||
-  ||| This is really hard however because all we have access to is the elaborated
+  ||| This is really hard however because all we have access to is the
   ||| clauses of the lifted case expression. So e.g.
   |||      f x = case g x of p -> e
   ||| became
@@ -63,21 +63,24 @@ mutual
   ||| name x.
   |||
   ||| We will try to do our best...
-  unelabCase : {auto c : Ref Ctxt Defs} ->
+  unelabCase : {vars : _} ->
+               {auto c : Ref Ctxt Defs} ->
                List (Name, Nat) ->
-               Name -> List IArg -> IRawImp ->
-               Core IRawImp
-  unelabCase nest n args orig
+               Env Term vars ->
+               Name ->
+               List (Term vars) ->
+               Core (Maybe IRawImp)
+  unelabCase nest env n args
       = do defs <- get Ctxt
            Just glob <- lookupCtxtExact n (gamma defs)
-                | Nothing => pure orig
+                | Nothing => pure Nothing
            let PMDef _ pargs treect _ pats = definition glob
-                | _ => pure orig
+                | _ => pure Nothing
            let Just argpos = findArgPos treect
-                | _ => pure orig
+                | _ => pure Nothing
            if length args == length pargs
-              then mkCase pats argpos 0 args
-              else pure orig
+              then mkCase pats argpos args
+              else pure Nothing
     where
       -- Need to find the position of the scrutinee to rebuild original
       -- case correctly
@@ -85,76 +88,84 @@ mutual
       findArgPos (Case idx p _ _) = Just idx
       findArgPos _ = Nothing
 
-      idxOrDefault : Nat -> a -> List a -> a
-      idxOrDefault Z def (x :: _) = x
-      idxOrDefault (S k) def (_ :: xs) = idxOrDefault k def xs
-      idxOrDefault _ def [] = def
+      idxOrMaybe : Nat -> List a -> Maybe a
+      idxOrMaybe Z (x :: _) = Just x
+      idxOrMaybe (S k) (_ :: xs) = idxOrMaybe k xs
+      idxOrMaybe _ [] = Nothing
 
-      getNth : Nat -> Term vars -> Term vars
-      getNth n tm
-          = case getFnArgs tm of
-                 (f, as) => idxOrDefault n f as
+      -- TODO: some utility like this should probably be implemented in Core
+      substVars : List (List (Var vs), Term vs) -> Term vs -> Term vs
+      substVars xs tm@(Local fc _ idx prf)
+          = case find (any ((idx ==) . varIdx) . fst) xs of
+                 Just (_, new) => new
+                 Nothing => tm
+      substVars xs (Meta fc n i args)
+          = Meta fc n i (map (substVars xs) args)
+      substVars xs (Bind fc y b scope)
+          = Bind fc y (map (substVars xs) b) (substVars (map (bimap (map weaken) weaken) xs) scope)
+      substVars xs (App fc fn arg)
+          = App fc (substVars xs fn) (substVars xs arg)
+      substVars xs (As fc s as pat)
+          = As fc s as (substVars xs pat)
+      substVars xs (TDelayed fc y z)
+          = TDelayed fc y (substVars xs z)
+      substVars xs (TDelay fc y t z)
+          = TDelay fc y (substVars xs t) (substVars xs z)
+      substVars xs (TForce fc r y)
+          = TForce fc r (substVars xs y)
+      substVars xs tm = tm
 
-      nthArg : FC -> Nat -> Term vars -> Term vars
-      nthArg fc drop (App afc f a) = getNth drop (App afc f a)
-      nthArg fc drop tm = Erased fc False
+      embedVar : Var vs -> Var (vs ++ more)
+      embedVar (MkVar p) = MkVar (varExtend p)
+
+      substArgs : SizeOf vs -> List (List (Var vs), Term vars) -> Term vs -> Term (vs ++ vars)
+      substArgs p substs tm =
+        let
+          substs' = map (bimap (map $ embedVar {more = vars}) (weakenNs p)) substs
+          tm' = embed tm
+        in
+          substVars substs' tm'
+
+      argVars : {vs : _} -> Term vs -> List (Var vs)
+      argVars (As _ _ as pat) = argVars as ++ argVars pat
+      argVars (Local _ _ _ p) = [MkVar p]
+      argVars _ = []
 
-      ||| This is where we should introduce some renaming in the RHS to match
-      ||| the specialised call.
       mkClause : FC -> Nat ->
+                 List (Term vars) ->
                  (vs ** (Env Term vs, Term vs, Term vs)) ->
-                 Core IImpClause
-      mkClause fc dropped (vs ** (env, lhs, rhs))
+                 Core (Maybe IImpClause)
+      mkClause fc argpos args (vs ** (clauseEnv, lhs, rhs))
           = do logTerm "unelab.case.clause" 20 "Unelaborating clause" lhs
-               let pat = nthArg fc dropped lhs
-               logTerm "unelab.case.clause" 20 "Unelaborating LHS pattern" pat
-               lhs' <- unelabTy Full nest env pat
+               let patArgs = snd (getFnArgs lhs)
+                   Just pat = idxOrMaybe argpos patArgs
+                     | _ => pure Nothing
+                   rhs = substArgs (mkSizeOf vs) (zip (map argVars patArgs) args) rhs
+               logTerm "unelab.case.clause" 20 "Unelaborating LHS" pat
+               lhs' <- unelabTy Full nest clauseEnv pat
                logTerm "unelab.case.clause" 20 "Unelaborating RHS" rhs
-               logEnv "unelab.case.clause" 20 "In Env" env
-               rhs' <- unelabTy Full nest env rhs
-               pure (PatClause fc (fst lhs') (fst rhs'))
+               logEnv "unelab.case.clause" 20 "In Env" clauseEnv
+               rhs' <- unelabTy Full nest (clauseEnv ++ env) rhs
+               pure $ Just $ (PatClause fc (fst lhs') (fst rhs'))
 
       ||| mkCase looks up the value passed as the scrutinee of the case-block.
-      ||| @ idx     is the running index of the case-block's scrutinee
-      |||           It starts as the actual index and is decreased as we pass
-      |||           arguments on the way to finding the scrutinee.
-      ||| @ dropped is the number of arguments already dropped
-      |||           It starts as 0 and increases as we pass arguments
-      ||| Invariant: idx + dropped = argpos
+      ||| @ argpos  is the index of the case-block's scrutinee in args
       ||| @ args    is the list of arguments at the call site of the case-block
       |||
       ||| Once we have the scrutinee `e`, we can form `case e of` and so focus
       ||| on manufacturing the clauses.
       mkCase : List (vs ** (Env Term vs, Term vs, Term vs)) ->
-               (idx : Nat) -> (droppped : Nat) -> (args : List IArg) ->
-               Core IRawImp
-      mkCase pats (S k) dropped (_ :: args) = mkCase pats k (S dropped) args
-      mkCase pats Z dropped (Explicit fc tm :: args)
+               (argpos : Nat) -> List (Term vars) -> Core (Maybe IRawImp)
+      mkCase pats argpos args
           = do unless (null args) $ log "unelab.case.clause" 20 $
                  unwords $ "Ignoring" :: map show args
-               pats' <- traverse (mkClause fc dropped) pats
-               pure $ ICase fc tm (Implicit fc False) pats'
-      mkCase _ _ _ _ = pure orig
-
-  unelabSugar : {auto c : Ref Ctxt Defs} ->
-                (umode : UnelabMode) ->
-                (nest : List (Name, Nat)) ->
-                (IRawImp, Glued vars) ->
-                Core (IRawImp, Glued vars)
-  unelabSugar (NoSugar _) nest res = pure res
-  unelabSugar ImplicitHoles nest res = pure res
-  unelabSugar _ nest (tm, ty)
-      = let (f, args) = getFnArgs tm [] in
-            case f of
-             IVar fc (MkKindedName _ _ (NS ns (CaseBlock n i))) =>
-               do log "unelab.case" 20 $ unlines
-                         [ "Unelaborating case " ++ show (n, i)
-                         , "with arguments: " ++ show args
-                         ]
-                  tm <- unelabCase nest (NS ns (CaseBlock n i)) args tm
-                  log "unelab.case" 20 $ "Unelaborated to: " ++ show tm
-                  pure (tm, ty)
-             _ => pure (tm, ty)
+               let Just scrutinee = idxOrMaybe argpos args
+                     | _ => pure Nothing
+                   fc = getLoc scrutinee
+               (tm, _) <- unelabTy Full nest env scrutinee
+               Just pats' <- map sequence $ traverse (mkClause fc argpos args) pats
+                 | _ => pure Nothing
+               pure $ Just $ ICase fc tm (Implicit fc False) pats'
 
   dropParams : {auto c : Ref Ctxt Defs} ->
                List (Name, Nat) -> (IRawImp, Glued vars) ->
@@ -185,7 +196,7 @@ mutual
              Env Term vars -> Term vars ->
              Core (IRawImp, Glued vars)
   unelabTy umode nest env tm
-      = unelabSugar umode nest !(dropParams nest !(unelabTy' umode nest env tm))
+      = dropParams nest !(unelabTy' umode nest env tm)
 
   unelabTy' : {vars : _} ->
               {auto c : Ref Ctxt Defs} ->
@@ -244,11 +255,23 @@ mutual
          = if n `elem` vs
               then uniqueLocal vs (next n)
               else n
-  unelabTy' umode nest env (App fc fn arg)
+  unelabTy' umode nest env tm@(App fc fn arg)
       = do (fn', gfnty) <- unelabTy umode nest env fn
            (arg', gargty) <- unelabTy umode nest env arg
            fnty <- getNF gfnty
            defs <- get Ctxt
+           Nothing <-
+              case umode of
+                (NoSugar _) => pure Nothing
+                ImplicitHoles => pure Nothing
+                _ => case getFnArgs tm of
+                     (Ref _ _ fnName, args) => do
+                       fullName <- getFullName fnName
+                       let (NS ns (CaseBlock n i)) = fullName
+                         | _ => pure Nothing
+                       unelabCase nest env fullName args
+                     _ => pure Nothing
+             | Just tm => pure (tm, gErased fc)
            case fnty of
                 NBind _ x (Pi _ rig Explicit ty) sc
                   => do sc' <- sc defs (toClosure defaultOpts env arg)

tests/Main.idr

@@ -169,7 +169,7 @@ idrisTestsBuiltin = MkTestPool "Builtin types and functions" [] Nothing
 idrisTestsEvaluator : TestPool
 idrisTestsEvaluator = MkTestPool "Evaluation" [] Nothing
       [ -- Evaluator
-       "evaluator001", "evaluator002", "evaluator003",
+       "evaluator001", "evaluator002", "evaluator003", "evaluator004",
        -- Miscellaneous REPL
        "interpreter001", "interpreter002", "interpreter003", "interpreter004",
        "interpreter005", "interpreter006", "interpreter007", "interpreter008"]

tests/idris2/debug001/expected

@@ -66,7 +66,7 @@ Main> Prelude.Types.SnocList.filter
 Arguments [{arg:0}, {arg:1}, {arg:2}]
 Compile time tree: case {arg:2} of
   Lin {e:0} => [<]
-  (:<) {e:1} {e:2} {e:3} => let rest = filter {arg:1} {e:2} in if {arg:1} {e:3} then rest :< x else rest
+  (:<) {e:1} {e:2} {e:3} => let rest = filter {arg:1} {e:2} in if {arg:1} {e:3} then rest :< {e:3} else rest
 Erasable args: [0]
 Inferrable args: [0]
 Compiled: \ {arg:1}, {arg:2} => case {arg:2} of

tests/idris2/evaluator004/Issue1282.idr

@@ -0,0 +1,10 @@
+-- https://github.com/idris-lang/Idris2/issues/1282#issue-852601557
+0 Alias : Type -> Type
+Alias a = (b : Bool) -> if b then a else a
+foo : Alias ()
+foo = ?test1
+
+-- https://github.com/idris-lang/Idris2/issues/2461#issue-1228334999
+test2 = \x => (\y => the Bool $ if y then y else y) x
+
+test_fold = \a, b, c => foldl (\acc, i => case i of Z => S acc; _ => acc) Z [a, b, c]

tests/idris2/evaluator004/expected

@@ -0,0 +1,5 @@
+1/1: Building Issue1282 (Issue1282.idr)
+Main> Main.test1 : (b : Bool) -> (if b then () else ())
+Main> \x => if x then x else x
+Main> \a, b, c => case c of { 0 => S (case b of { 0 => S (case a of { 0 => 1 ; i => 0 }) ; i => case a of { 0 => 1 ; i => 0 } }) ; i => case b of { 0 => S (case a of { 0 => 1 ; i => 0 }) ; i => case a of { 0 => 1 ; i => 0 } } }
+Main> Bye for now!

tests/idris2/evaluator004/input

@@ -0,0 +1,4 @@
+:t test1
+test2
+test_fold
+:q

tests/idris2/evaluator004/run

@@ -0,0 +1,3 @@
+rm -rf build
+
+$1 --no-banner --no-color --console-width 0 Issue1282.idr < input

tests/idris2/interactive017/expected

@@ -1,4 +1,4 @@
 1/1: Building RLE (RLE.idr)
 RLE> uncompress Empty = Val []
-uncompress (Run n x rest) = let Val ys = uncompress rest in Val (rep n x ++ ys)
+uncompress (Run n x rest) = let Val ys = uncompress rest in Val (rep n x ++ _)
 RLE> Bye for now!