Implement local function definitions

Changed nested names (and case blocks) to store the resolved name as the
outer name, rather than the unresolved name, otherwise we'll have issues
when loading from TTC

src/Core/Name.idr

@@ -8,9 +8,9 @@ data Name : Type where
      UN : String -> Name -- user defined name
      MN : String -> Int -> Name -- machine generated name
      PV : Name -> Int -> Name -- pattern variable name; int is the resolved function id
-     Nested : Name -> Name -> Name -- nested function name
-     CaseBlock : Name -> Int -> Name -- case block nested in name
-     WithBlock : Name -> Int -> Name -- with block nested in name
+     Nested : Int -> Name -> Name -- nested function name
+     CaseBlock : Int -> Int -> Name -- case block nested in (resolved) name
+     WithBlock : Int -> Int -> Name -- with block nested in (resolved) name
      Resolved : Int -> Name -- resolved, index into context
 
 export
@@ -32,8 +32,8 @@ nameRoot (UN n) = n
 nameRoot (MN n _) = n
 nameRoot (PV n _) = nameRoot n
 nameRoot (Nested n inner) = nameRoot inner
-nameRoot (CaseBlock n inner) = nameRoot n
-nameRoot (WithBlock n inner) = nameRoot n
+nameRoot (CaseBlock n inner) = "$" ++ show n
+nameRoot (WithBlock n inner) = "$" ++ show n
 nameRoot (Resolved i) = "$" ++ show i
 
 --- Drop a namespace from a name
@@ -139,21 +139,21 @@ nameEq (PV x t) (PV y t') with (nameEq x y)
     nameEq (PV y t) (PV y t) | (Just Refl) | (Yes Refl) = Just Refl
     nameEq (PV y t) (PV y t') | (Just Refl) | (No p) = Nothing
   nameEq (PV x t) (PV y t') | Nothing = Nothing
-nameEq (Nested x y) (Nested x' y') with (nameEq x x')
-  nameEq (Nested x y) (Nested x' y') | Nothing = Nothing
-  nameEq (Nested x y) (Nested x y') | (Just Refl) with (nameEq y y')
-    nameEq (Nested x y) (Nested x y') | (Just Refl) | Nothing = Nothing
-    nameEq (Nested x y) (Nested x y) | (Just Refl) | (Just Refl) = Just Refl
-nameEq (CaseBlock x y) (CaseBlock x' y') with (nameEq x x')
-  nameEq (CaseBlock x y) (CaseBlock x' y') | Nothing = Nothing
-  nameEq (CaseBlock x y) (CaseBlock x y') | (Just Refl) with (decEq y y')
-    nameEq (CaseBlock x y) (CaseBlock x y') | (Just Refl) | (No p) = Nothing
-    nameEq (CaseBlock x y) (CaseBlock x y) | (Just Refl) | (Yes Refl) = Just Refl
-nameEq (WithBlock x y) (WithBlock x' y') with (nameEq x x')
-  nameEq (WithBlock x y) (WithBlock x' y') | Nothing = Nothing
-  nameEq (WithBlock x y) (WithBlock x y') | (Just Refl) with (decEq y y')
-    nameEq (WithBlock x y) (WithBlock x y') | (Just Refl) | (No p) = Nothing
-    nameEq (WithBlock x y) (WithBlock x y) | (Just Refl) | (Yes Refl) = Just Refl
+nameEq (Nested x y) (Nested x' y') with (decEq x x')
+  nameEq (Nested x y) (Nested x' y') | (No p) = Nothing
+  nameEq (Nested x y) (Nested x y') | (Yes Refl) with (nameEq y y')
+    nameEq (Nested x y) (Nested x y') | (Yes Refl) | Nothing = Nothing
+    nameEq (Nested x y) (Nested x y) | (Yes Refl) | (Just Refl) = Just Refl
+nameEq (CaseBlock x y) (CaseBlock x' y') with (decEq x x')
+  nameEq (CaseBlock x y) (CaseBlock x' y') | (No p) = Nothing
+  nameEq (CaseBlock x y) (CaseBlock x y') | (Yes Refl) with (decEq y y')
+    nameEq (CaseBlock x y) (CaseBlock x y') | (Yes Refl) | (No p) = Nothing
+    nameEq (CaseBlock x y) (CaseBlock x y) | (Yes Refl) | (Yes Refl) = Just Refl
+nameEq (WithBlock x y) (WithBlock x' y') with (decEq x x')
+  nameEq (WithBlock x y) (WithBlock x' y') | (No p) = Nothing
+  nameEq (WithBlock x y) (WithBlock x y') | (Yes Refl) with (decEq y y')
+    nameEq (WithBlock x y) (WithBlock x y') | (Yes Refl) | (No p) = Nothing
+    nameEq (WithBlock x y) (WithBlock x y) | (Yes Refl) | (Yes Refl) = Just Refl
 nameEq (Resolved x) (Resolved y) with (decEq x y)
   nameEq (Resolved y) (Resolved y) | (Yes Refl) = Just Refl
   nameEq (Resolved x) (Resolved y) | (No contra) = Nothing

src/Core/TTC.idr

@@ -4,6 +4,7 @@ import Core.CaseTree
 import Core.Core
 import Core.Env
 import Core.FC
+import Core.Name
 import Core.TT
 
 import Utils.Binary

src/Core/Unify.idr

@@ -247,18 +247,18 @@ instantiate : {auto c : Ref Ctxt Defs} ->
 instantiate {newvars} loc env mname mref mdef locs otm tm
     = do log 5 $ "Instantiating " ++ show tm ++ " in " ++ show newvars 
          let ty = type mdef
-         case mkDef [] newvars (snocList newvars) CompatPre
+         defs <- get Ctxt
+         rhs <- mkDef [] newvars (snocList newvars) CompatPre
                      (rewrite appendNilRightNeutral newvars in locs)
-                     (rewrite appendNilRightNeutral newvars in tm) ty of
-               Nothing => ufail loc $ "Can't make solution for " ++ show mname
-               Just rhs =>
-                  do logTerm 5 ("Instantiated: " ++ show mname) ty
-                     logTerm 5 "Definition" rhs
-                     let newdef = record { definition = 
-                                             PMDef [] (STerm rhs) (STerm rhs) [] 
-                                         } mdef
-                     addDef (Resolved mref) newdef
-                     removeHole mref
+                     (rewrite appendNilRightNeutral newvars in tm) 
+                     !(nf defs [] ty)
+         logTerm 5 ("Instantiated: " ++ show mname) ty
+         logTerm 5 "Definition" rhs
+         let newdef = record { definition = 
+                                 PMDef [] (STerm rhs) (STerm rhs) [] 
+                             } mdef
+         addDef (Resolved mref) newdef
+         removeHole mref
   where
     updateLoc : {v : Nat} -> List (Var vs) -> .(IsVar name v vs') -> 
                 Maybe (Var vs)
@@ -295,18 +295,22 @@ instantiate {newvars} loc env mname mref mdef locs otm tm
     mkDef : (got : List Name) -> (vs : List Name) -> SnocList vs ->
             CompatibleVars got rest ->
             List (Var (vs ++ got)) -> Term (vs ++ got) -> 
-            Term rest -> Maybe (Term rest)
+            NF [] -> Core (Term rest)
     mkDef got [] Empty cvs locs tm ty 
-        = do tm' <- updateLocs (reverse locs) tm
+        = do let Just tm' = updateLocs (reverse locs) tm
+                    | Nothing => ufail loc ("Can't make solution for " ++ show mname)
              pure (renameVars cvs tm')
-    mkDef got (vs ++ [v]) (Snoc rec) cvs locs tm (Bind bfc x (Pi c _ ty) sc) 
-        = do sc' <- mkDef (v :: got) vs rec (CompatExt cvs)
+    mkDef got (vs ++ [v]) (Snoc rec) cvs locs tm (NBind bfc x (Pi c _ ty) scfn) 
+        = do defs <- get Ctxt
+             sc <- scfn defs (toClosure defaultOpts [] (Erased bfc))
+             sc' <- mkDef (v :: got) vs rec (CompatExt cvs)
                        (rewrite appendAssociative vs [v] got in locs)
                        (rewrite appendAssociative vs [v] got in tm)
                        sc
              pure (Bind bfc x (Lam c Explicit (Erased bfc)) sc')
-    mkDef got (vs ++ [v]) (Snoc rec) cvs locs tm ty = Nothing
-    
+    mkDef got (vs ++ [v]) (Snoc rec) cvs locs tm ty
+        = ufail loc $ "Can't make solution for " ++ show mname  
+
 isDefInvertible : {auto c : Ref Ctxt Defs} ->
                   Int -> Core Bool
 isDefInvertible i

src/Core/UnifyState.idr

@@ -256,6 +256,12 @@ applyTo {vars} fc tm env
   = let args = reverse (mkConstantAppArgs {done = []} False fc env []) in
         apply fc (explApp Nothing) tm (rewrite sym (appendNilRightNeutral vars) in args)
 
+export
+applyToFull : FC -> Term vars -> Env Term vars -> Term vars
+applyToFull {vars} fc tm env
+  = let args = reverse (mkConstantAppArgs {done = []} True fc env []) in
+        apply fc (explApp Nothing) tm (rewrite sym (appendNilRightNeutral vars) in args)
+
 -- Create a new metavariable with the given name and return type,
 -- and return a term which is the metavariable applied to the environment
 -- (and which has the given type)

src/TTImp/Elab/App.idr

@@ -14,12 +14,15 @@ import TTImp.TTImp
 
 %default covering
 
+-- Get the type of a variable, assuming we haven't found it in the nested
+-- names. Look in the Env first, then the global context.
 getNameType : {vars : _} ->
               {auto c : Ref Ctxt Defs} ->
               {auto e : Ref EST (EState vars)} ->
-              RigCount -> Env Term vars -> FC -> Name -> Maybe (Glued vars) ->
+              RigCount -> Env Term vars ->
+              FC -> Name ->
               Core (Term vars, Glued vars)
-getNameType rigc env fc x expected
+getNameType rigc env fc x
     = case defined x env of
            Just (MkIsDefined rigb lv) => 
               do rigSafe rigb rigc
@@ -43,6 +46,38 @@ getNameType rigc env fc x expected
     rigSafe Rig0 Rig1 = throw (LinearMisuse fc x Rig0 Rig1)
     rigSafe _ _ = pure ()
 
+-- Get the type of a variable, looking it up in the nested names first.
+getVarType : {vars : _} ->
+             {auto c : Ref Ctxt Defs} ->
+             {auto e : Ref EST (EState vars)} ->
+             RigCount -> NestedNames vars -> Env Term vars ->
+             FC -> Name ->
+             Core (Term vars, Glued vars)
+getVarType rigc nest env fc x
+    = case lookup x (names nest) of
+           Nothing => getNameType rigc env fc x
+           Just (nestn, tmf) =>
+              do defs <- get Ctxt
+                 let n' = maybe x id nestn
+                 case !(lookupCtxtExact n' (gamma defs)) of
+                      Nothing => throw (UndefinedName fc n')
+                      Just ndef =>
+                         let nt = case definition ndef of
+                                       PMDef _ _ _ _ => Func
+                                       DCon t a => DataCon t a
+                                       TCon t a _ _ _ _ => TyCon t a
+                                       _ => Func
+                             tm = tmf fc nt
+                             tyenv = useVars (map snd (getArgs tm))
+                                             (embed (type ndef)) in
+                             pure (tm, gnf env tyenv)
+    where
+      useVars : List (Term vars) -> Term vars -> Term vars
+      useVars [] sc = sc
+      useVars (a :: as) (Bind bfc n (Pi c _ ty) sc) 
+           = Bind bfc n (Let c a ty) (useVars (map weaken as) sc)
+      useVars _ sc = sc -- Can't happen?
+
 isHole : NF vars -> Bool
 isHole (NApp _ (NMeta _ _ _) _) = True
 isHole _ = False
@@ -304,7 +339,7 @@ checkApp rig elabinfo nest env fc (IApp fc' fn arg) expargs impargs exp
 checkApp rig elabinfo nest env fc (IImplicitApp fc' fn nm arg) expargs impargs exp
    = checkApp rig elabinfo nest env fc' fn expargs ((nm, arg) :: impargs) exp
 checkApp rig elabinfo nest env fc (IVar fc' n) expargs impargs exp
-   = do (ntm, nty_in) <- getNameType rig env fc n Nothing
+   = do (ntm, nty_in) <- getVarType rig nest env fc n
         nty <- getNF nty_in
         logC 10 (do defs <- get Ctxt
                     fnty <- quote defs env nty

src/TTImp/Elab/Local.idr

@@ -0,0 +1,99 @@
+module TTImp.Elab.Local
+
+import Core.Context
+import Core.Core
+import Core.Env
+import Core.Normalise
+import Core.Unify
+import Core.TT
+import Core.Value
+
+import TTImp.Elab.Check
+import TTImp.TTImp
+
+%default covering
+
+definedInBlock : List ImpDecl -> List Name
+definedInBlock = concatMap defName
+  where
+    getName : ImpTy -> Name
+    getName (MkImpTy _ n _) = n
+
+    defName : ImpDecl -> List Name
+    defName (IClaim _ _ _ _ ty) = [getName ty]
+    defName (IData _ _ (MkImpData _ n _ _ cons)) = n :: map getName cons
+    defName (IData _ _ (MkImpLater _ n _)) = [n]
+-- TODO after adding these cases to ImpDecl
+--     defName (IParameters _ _ pds) = concatMap defName pds
+--     defName (IRecord _ _ (MkImpRecord _ n _ _ _)) = [n]
+    defName _ = []
+
+export
+checkLocal : {vars : _} ->
+             {auto c : Ref Ctxt Defs} ->
+             {auto u : Ref UST UState} ->
+             {auto e : Ref EST (EState vars)} ->
+             RigCount -> ElabInfo -> 
+             NestedNames vars -> Env Term vars -> 
+             FC -> List ImpDecl -> (scope : RawImp) ->
+             (expTy : Maybe (Glued vars)) ->
+             Core (Term vars, Glued vars)
+checkLocal {vars} rig elabinfo nest env fc nestdecls scope expty
+    = do let defNames = definedInBlock nestdecls
+         est <- get EST
+         let f = defining est
+         names' <- traverse (applyEnv f) defNames
+         let nest' = record { names $= (names' ++) } nest
+         let env' = dropLinear env
+         traverse (processDecl nest' env') (map (updateName nest') nestdecls)
+         check rig elabinfo nest' env scope expty
+  where
+      -- For the local definitions, don't allow access to linear things
+      -- unless they're explicitly passed.
+      -- This is because, at the moment, we don't have any mechanism of
+      -- ensuring the nested definition is used exactly once
+      dropLinear : Env Term vs -> Env Term vs
+      dropLinear [] = []
+      dropLinear (b :: bs) 
+          = if isLinear (multiplicity b)
+               then setMultiplicity b Rig0 :: dropLinear bs
+               else b :: dropLinear bs
+
+      applyEnv : Int -> Name -> 
+                 Core (Name, (Maybe Name, FC -> NameType -> Term vars))
+      applyEnv outer inner 
+            = do n' <- resolveName (Nested outer inner)
+                 pure (inner, (Just (Resolved n'), 
+                          \fc, nt => applyToFull fc 
+                                 (Ref fc nt (Resolved n')) env))
+
+      -- Update the names in the declarations to the new 'nested' names.
+      -- When we encounter the names in elaboration, we'll update to an
+      -- application of the nested name.
+      newName : NestedNames vars -> Name -> Name
+      newName nest n 
+          = case lookup n (names nest) of
+                 Just (Just n', _) => n'
+                 _ => n
+
+      updateTyName : NestedNames vars -> ImpTy -> ImpTy
+      updateTyName nest (MkImpTy loc' n ty) 
+          = MkImpTy loc' (newName nest n) ty
+
+      updateDataName : NestedNames vars -> ImpData -> ImpData
+      updateDataName nest (MkImpData loc' n tycons dopts dcons)
+          = MkImpData loc' (newName nest n) tycons dopts
+                           (map (updateTyName nest) dcons)
+      updateDataName nest (MkImpLater loc' n tycons)
+          = MkImpLater loc' (newName nest n) tycons
+
+      updateName : NestedNames vars -> ImpDecl -> ImpDecl
+      updateName nest (IClaim loc' r vis fnopts ty) 
+           = IClaim loc' r vis fnopts (updateTyName nest ty)
+      updateName nest (IDef loc' n cs) 
+           = IDef loc' (newName nest n) cs
+      updateName nest (IData loc' vis d) 
+           = IData loc' vis (updateDataName nest d)
+      updateName nest i = i
+
+

src/TTImp/Elab/Term.idr

@@ -14,6 +14,7 @@ import TTImp.Elab.As
 import TTImp.Elab.Binders
 import TTImp.Elab.Check
 import TTImp.Elab.ImplicitBind
+import TTImp.Elab.Local
 import TTImp.Elab.Prim
 import TTImp.TTImp
 
@@ -98,7 +99,7 @@ checkTerm rig elabinfo nest env (ILet fc r n nTy nVal scope) exp
 checkTerm rig elabinfo nest env (ICase fc scr scrty als) exp
     = throw (InternalError "case not implemented")
 checkTerm rig elabinfo nest env (ILocal fc nested scope) exp
-    = throw (InternalError "let functions not implemented")
+    = checkLocal rig elabinfo nest env fc nested scope exp
 checkTerm rig elabinfo nest env (IUpdate fc upds rec) exp
     = throw (InternalError "record update not implemented")
 checkTerm rig elabinfo nest env (IApp fc fn arg) exp 

src/TTImp/Parser.idr

@@ -374,7 +374,7 @@ directive fname indents
          s <- name
          end <- location
          let fc = MkFC fname start end
-         pure (IPragma (\c, env => setPair {c} fc p f s))
+         pure (IPragma (\c, nest, env => setPair {c} fc p f s))
 
 -- Declared at the top
 -- topDecl : FileName -> IndentInfo -> Rule ImpDecl

src/TTImp/ProcessDecls.idr

@@ -31,7 +31,7 @@ process nest env (INamespace fc ns decls)
          traverse (processDecl nest env) decls
          setNS oldns
 process {c} nest env (IPragma act)
-    = act c env
+    = act c nest env
 process nest env (ILog n)
     = setLogLevel n
 

src/TTImp/ProcessDef.idr

@@ -136,7 +136,8 @@ checkClause : {auto c : Ref Ctxt Defs} ->
 checkClause mult hashit n nest env (ImpossibleClause fc lhs)
     = throw (InternalError "impossible not implemented yet")
 checkClause mult hashit n nest env (PatClause fc lhs_in rhs)
-    = do lhs <- lhsInCurrentNS lhs_in 
+    = do lhs <- lhsInCurrentNS nest lhs_in 
+         log 5 $ "Checking " ++ show lhs
          (lhstm, lhstyg) <- elabTerm n (InLHS mult) nest env 
                                 (IBindHere fc PATTERN lhs) Nothing
          lhsty <- getTerm lhstyg

src/TTImp/TTImp.idr

@@ -12,15 +12,17 @@ record NestedNames (vars : List Name) where
   constructor MkNested
   -- A map from names to the decorated version of the name, and the new name
   -- applied to its enclosing environment
-  names : List (Name, (Maybe Name, Term vars))
+  -- Takes the location and name type, because we don't know them until we
+  -- elaborate the name at the point of use
+  names : List (Name, (Maybe Name, FC -> NameType -> Term vars))
 
 export
 Weaken NestedNames where
   weaken (MkNested ns) = MkNested (map wknName ns)
     where
-      wknName : (Name, (Maybe Name, Term vars)) ->
-                (Name, (Maybe Name, Term (n :: vars)))
-      wknName (n, (mn, rep)) = (n, (mn, weaken rep))
+      wknName : (Name, (Maybe Name, FC -> NameType -> Term vars)) ->
+                (Name, (Maybe Name, FC -> NameType -> Term (n :: vars)))
+      wknName (n, (mn, rep)) = (n, (mn, \fc, nt => weaken (rep fc nt)))
 
 -- Unchecked terms, with implicit arguments
 -- This is the raw, elaboratable form.
@@ -194,7 +196,8 @@ mutual
        IData : FC -> Visibility -> ImpData -> ImpDecl
        IDef : FC -> Name -> List ImpClause -> ImpDecl
        INamespace : FC -> List String -> List ImpDecl -> ImpDecl 
-       IPragma : ({vars : _} -> Ref Ctxt Defs -> Env Term vars -> Core ()) -> 
+       IPragma : ({vars : _} -> Ref Ctxt Defs -> 
+                  NestedNames vars -> Env Term vars -> Core ()) -> 
                  ImpDecl
        ILog : Nat -> ImpDecl
 
@@ -220,23 +223,21 @@ data ImpREPL : Type where
 
 export
 lhsInCurrentNS : {auto c : Ref Ctxt Defs} ->
-                 RawImp -> Core RawImp
-lhsInCurrentNS (IApp loc f a)
-    = do f' <- lhsInCurrentNS f
+                 NestedNames vars -> RawImp -> Core RawImp
+lhsInCurrentNS nest (IApp loc f a)
+    = do f' <- lhsInCurrentNS nest f
          pure (IApp loc f' a)
-lhsInCurrentNS (IImplicitApp loc f n a)
-    = do f' <- lhsInCurrentNS f
+lhsInCurrentNS nest (IImplicitApp loc f n a)
+    = do f' <- lhsInCurrentNS nest f
          pure (IImplicitApp loc f' n a)
-lhsInCurrentNS tm@(IVar loc (NS _ _)) = pure tm -- leave explicit NS alone
-lhsInCurrentNS (IVar loc n)
-       = do n' <- inCurrentNS n
-            pure (IVar loc n')
---     = case lookup n (names nest) of
---            Nothing =>
---               do n' <- inCurrentNS n
---                  pure (IVar loc n')
---            -- If it's one of the names in the current nested block, we'll
---            -- be rewriting it during elaboration to be in the scope of the
---            -- parent name.
---            Just _ => pure (IVar loc n)
-lhsInCurrentNS tm = pure tm
+lhsInCurrentNS nest tm@(IVar loc (NS _ _)) = pure tm -- leave explicit NS alone
+lhsInCurrentNS nest (IVar loc n)
+    = case lookup n (names nest) of
+           Nothing =>
+              do n' <- inCurrentNS n
+                 pure (IVar loc n')
+           -- If it's one of the names in the current nested block, we'll
+           -- be rewriting it during elaboration to be in the scope of the
+           -- parent name.
+           Just _ => pure (IVar loc n)
+lhsInCurrentNS nest tm = pure tm

tests/ttimp/nest001/Let.yaff

@@ -1,3 +1,42 @@
 testlet : Integer -> Integer
 testlet $x = let y = prim__add_Integer x x in 
                  prim__add_Integer y y 
+
+data Nat : Type where
+     Z : Nat
+     S : Nat -> Nat
+
+plus : Nat -> Nat -> Nat
+plus Z     $y = y
+plus (S $k) $y = S (plus k y)
+
+fn2 : Nat -> Nat -> Nat
+fn2 $x $y 
+    = let w : Nat
+          w = plus x x;
+          help : Nat -> Nat
+          help $z = plus z w in
+          plus (help y) w
+
+fn : Nat -> Nat -> Nat
+fn $x $y 
+    = let w : ?; w = plus x x in
+      let foo = plus x x in
+      let help : Nat -> Nat;
+          help $z = plus z w in
+      plus (help y) foo
+
+localdata : Nat -> Nat
+localdata $var =
+   let data Bool : Type where
+            False : Bool
+            True : Bool
+
+       isS : Nat -> Bool
+       isS Z = False
+       isS (S $k) = True
+
+       boolToNat : Bool -> Nat
+       boolToNat False = Z
+       boolToNat True = S var in
+   boolToNat (isS var)

tests/ttimp/nest001/expected

@@ -1,4 +1,8 @@
 Processing as TTImp
 Written TTC
 Yaffle> 8
+Yaffle> (Main.S (Main.S (Main.S (Main.S (Main.S Main.Z)))))
+Yaffle> (Main.S (Main.S (Main.S (Main.S (Main.S Main.Z)))))
+Yaffle> Main.Z
+Yaffle> (Main.S (Main.S (Main.S Main.Z)))
 Yaffle> Bye for now!

tests/ttimp/nest001/input

@@ -1,2 +1,6 @@
 testlet 2
+fn (S Z) (S Z)
+fn2 (S Z) (S Z)
+localdata Z
+localdata (S (S Z))
 :q

yaffle.ipkg

@@ -46,6 +46,7 @@ modules =
     TTImp.Elab.Binders,
     TTImp.Elab.Check,
     TTImp.Elab.ImplicitBind,
+    TTImp.Elab.Local,
     TTImp.Elab.Prim,
     TTImp.Elab.Term,
     TTImp.Parser,