Speed up comparison of environments

CTT.hs

@@ -36,8 +36,10 @@ data Branch = OBranch LIdent [Ident] Ter
   deriving (Eq,Show)
 
 -- Declarations: x : A = e
+-- A group of mutual declarations is identified by its location. It is used to
+-- speed up the Eq instance for Ctxt.
 type Decl  = (Ident,(Ter,Ter))
-data Decls = MutualDecls [Decl]
+data Decls = MutualDecls Loc [Decl]
            | OpaqueDecl Ident
            | TransparentDecl Ident
            | TransparentAllDecl
@@ -222,8 +224,18 @@ constPath = VPLam (Name "_")
 data Ctxt = Empty
           | Upd Ident Ctxt
           | Sub Name Ctxt
-          | Def [Decl] Ctxt
-  deriving (Show,Eq)
+          | Def Loc [Decl] Ctxt
+  deriving (Show)
+
+instance Eq Ctxt where
+    c == d = case (c, d) of
+        (Empty, Empty)              -> True
+        (Upd x c', Upd y d')        -> x == y && c' == d'
+        (Sub i c', Sub j d')        -> i == j && c' == d'
+        (Def m xs c', Def n ys d')  -> (m == n || xs == ys) && c' == d'
+            -- Invariant: if two declaration groups come from the same
+            -- location, they are equal and their contents are not compared.
+        _                           -> False
 
 -- The Idents and Names in the Ctxt refer to the elements in the two
 -- lists. This is more efficient because acting on an environment now
@@ -235,13 +247,13 @@ emptyEnv :: Env
 emptyEnv = (Empty,[],[],Nameless Set.empty)
 
 def :: Decls -> Env -> Env
-def (MutualDecls ds) (rho,vs,fs,Nameless os) = (Def ds rho,vs,fs,Nameless (os Set.\\ Set.fromList (declIdents ds)))
+def (MutualDecls m ds) (rho,vs,fs,Nameless os) = (Def m ds rho,vs,fs,Nameless (os Set.\\ Set.fromList (declIdents ds)))
 def (OpaqueDecl n) (rho,vs,fs,Nameless os) = (rho,vs,fs,Nameless (Set.insert n os))
 def (TransparentDecl n) (rho,vs,fs,Nameless os) = (rho,vs,fs,Nameless (Set.delete n os))
 def TransparentAllDecl (rho,vs,fs,Nameless os) = (rho,vs,fs,Nameless Set.empty)
 
 defWhere :: Decls -> Env -> Env
-defWhere (MutualDecls ds) (rho,vs,fs,Nameless os) = (Def ds rho,vs,fs,Nameless (os Set.\\ Set.fromList (declIdents ds)))
+defWhere (MutualDecls m ds) (rho,vs,fs,Nameless os) = (Def m ds rho,vs,fs,Nameless (os Set.\\ Set.fromList (declIdents ds)))
 defWhere (OpaqueDecl _) rho = rho
 defWhere (TransparentDecl _) rho = rho
 defWhere TransparentAllDecl rho = rho
@@ -276,18 +288,18 @@ formulaOfEnv = snd . valAndFormulaOfEnv
 domainEnv :: Env -> [Name]
 domainEnv (rho,_,_,_) = domCtxt rho
   where domCtxt rho = case rho of
-          Empty    -> []
-          Upd _ e  -> domCtxt e
-          Def ts e -> domCtxt e
-          Sub i e  -> i : domCtxt e
+          Empty      -> []
+          Upd _ e    -> domCtxt e
+          Def _ ts e -> domCtxt e
+          Sub i e    -> i : domCtxt e
 
 -- Extract the context from the environment, used when printing holes
 contextOfEnv :: Env -> [String]
 contextOfEnv rho = case rho of
   (Empty,_,_,_)               -> []
   (Upd x e,VVar n t:vs,fs,os) -> (n ++ " : " ++ show t) : contextOfEnv (e,vs,fs,os)
   (Upd x e,v:vs,fs,os)        -> (x ++ " = " ++ show v) : contextOfEnv (e,vs,fs,os)
-  (Def _ e,vs,fs,os)          -> contextOfEnv (e,vs,fs,os)
+  (Def _ _ e,vs,fs,os)        -> contextOfEnv (e,vs,fs,os)
   (Sub i e,vs,phi:fs,os)      -> (show i ++ " = " ++ show phi) : contextOfEnv (e,vs,fs,os)
 
 --------------------------------------------------------------------------------
@@ -307,11 +319,11 @@ showEnv b e =
           showEnv1 (env,us,fs,os) <+> names x <+> showVal1 u <> com
         (Sub i env,us,phi:fs,os) ->
           showEnv1 (env,us,fs,os) <+> names (show i) <+> text (show phi) <> com
-        (Def _ env,vs,fs,os)     -> showEnv1 (env,vs,fs,os)
+        (Def _ _ env,vs,fs,os)   -> showEnv1 (env,vs,fs,os)
         _                        -> showEnv b e
   in case e of
     (Empty,_,_,_)            -> PP.empty
-    (Def _ env,vs,fs,os)     -> showEnv b (env,vs,fs,os)
+    (Def _ _ env,vs,fs,os)   -> showEnv b (env,vs,fs,os)
     (Upd x env,u:us,fs,os)   ->
       par $ showEnv1 (env,us,fs,os) <+> names x <+> showVal1 u
     (Sub i env,us,phi:fs,os) ->
@@ -380,7 +392,7 @@ showTer1 t = case t of
   _        -> parens (showTer t)
 
 showDecls :: Decls -> Doc
-showDecls (MutualDecls defs) =
+showDecls (MutualDecls _ defs) =
   hsep $ punctuate comma
   [ text x <+> equals <+> showTer d | (x,(_,d)) <- defs ]
 showDecls (OpaqueDecl i) = text "opaque" <+> text i

Eval.hs

@@ -7,6 +7,7 @@ import Data.Map (Map,(!),mapWithKey,assocs,filterWithKey
                 ,elems,intersectionWith,intersection,keys
                 ,member,notMember,empty)
 import qualified Data.Map as Map
+import qualified Data.Set as Set
 
 import Connections
 import CTT
@@ -17,7 +18,7 @@ import CTT
 look :: String -> Env -> Val
 look x (Upd y rho,v:vs,fs,os) | x == y = v
                               | otherwise = look x (rho,vs,fs,os)
-look x r@(Def decls rho,vs,fs,Nameless os) = case lookup x decls of
+look x r@(Def _ decls rho,vs,fs,Nameless os) = case lookup x decls of
   Just (_,t) -> eval r t
   Nothing    -> look x (rho,vs,fs,Nameless os)
 look x (Sub _ rho,vs,_:fs,os) = look x (rho,vs,fs,os)
@@ -28,15 +29,15 @@ lookType x (Upd y rho,v:vs,fs,os)
   | x /= y        = lookType x (rho,vs,fs,os)
   | VVar _ a <- v = a
   | otherwise     = error ""
-lookType x r@(Def decls rho,vs,fs,os) = case lookup x decls of
+lookType x r@(Def _ decls rho,vs,fs,os) = case lookup x decls of
   Just (a,_) -> eval r a
   Nothing -> lookType x (rho,vs,fs,os)
 lookType x (Sub _ rho,vs,_:fs,os) = lookType x (rho,vs,fs,os)
 lookType x (Empty,_,_,_)                  = error $ "lookType: not found " ++ show x
 
 lookName :: Name -> Env -> Formula
 lookName i (Upd _ rho,v:vs,fs,os) = lookName i (rho,vs,fs,os)
-lookName i (Def _ rho,vs,fs,os)   = lookName i (rho,vs,fs,os)
+lookName i (Def _ _ rho,vs,fs,os) = lookName i (rho,vs,fs,os)
 lookName i (Sub j rho,vs,phi:fs,os) | i == j    = phi
                                     | otherwise = lookName i (rho,vs,fs,os)
 lookName i _ = error $ "lookName: not found " ++ show i
@@ -150,7 +151,7 @@ eval rho@(_,_,_,Nameless os) v = case v of
   U                   -> VU
   App r s             -> app (eval rho r) (eval rho s)
   Var i
-    | i `elem` os     -> VOpaque i (lookType i rho)
+    | i `Set.member` os -> VOpaque i (lookType i rho)
     | otherwise       -> look i rho
   Pi t@(Lam _ a _)    -> VPi (eval rho a) (eval rho t)
   Sigma t@(Lam _ a _) -> VSigma (eval rho a) (eval rho t)

Resolver.hs

@@ -8,6 +8,7 @@ import Control.Monad
 import Control.Monad.Reader
 import Control.Monad.Except
 import Control.Monad.Identity
+import Data.Maybe
 import Data.List
 import Data.Map (Map,(!))
 import qualified Data.Map as Map
@@ -332,6 +333,22 @@ resolveRTele (i:is) (t:ts) = do
   as <- local (insertVar i) (resolveRTele is ts)
   return ((i,a):as)
 
+-- Best effort to find the location of a declaration. This implementation
+-- returns the location of the first identifier it contains.
+findDeclLoc :: Decl -> Resolver Loc
+findDeclLoc d = getLoc loc
+    where loc = fromMaybe (-1, 0) $ mloc d
+          mloc d = case d of
+            DeclDef (AIdent (l, _)) _ _ _   -> Just l
+            DeclData (AIdent (l, _)) _ _    -> Just l
+            DeclHData (AIdent (l, _)) _ _   -> Just l
+            DeclSplit (AIdent (l, _)) _ _ _ -> Just l
+            DeclUndef (AIdent (l, _)) _ _   -> Just l
+            DeclMutual ds                   -> listToMaybe $ mapMaybe mloc ds
+            DeclOpaque (AIdent (l, _))      -> Just l
+            DeclTransparent (AIdent (l, _)) -> Just l
+            DeclTransparentAll              -> Nothing
+
 -- Resolve a declaration
 resolveDecl :: Decl -> Resolver (CTT.Decls,[(Ident,SymKind)])
 resolveDecl d = case d of
@@ -345,7 +362,8 @@ resolveDecl d = case d of
     -- mutual block
     ds <- sequence $ map (local (insertIdents ns)) bs
     let ads = zipWith (\ (x,y) z -> (x,(y,z))) as ds
-    return (CTT.MutualDecls ads,ns)
+    l <- findDeclLoc d
+    return (CTT.MutualDecls l ads,ns)
   DeclOpaque i  -> do
     resolveVar i
     return (CTT.OpaqueDecl (unAIdent i), [])
@@ -354,9 +372,10 @@ resolveDecl d = case d of
     return (CTT.TransparentDecl (unAIdent i), [])
   DeclTransparentAll -> return (CTT.TransparentAllDecl, [])
   _ -> do let (f,typ,body,ns) = resolveNonMutualDecl d
+          l <- findDeclLoc d
           a <- typ
           d <- body
-          return (CTT.MutualDecls [(f,(a,d))],ns)
+          return (CTT.MutualDecls l [(f,(a,d))],ns)
 
 resolveDecls :: [Decl] -> Resolver ([CTT.Decls],[(Ident,SymKind)])
 resolveDecls []     = return ([],[])

TypeChecker.hs

@@ -222,14 +222,14 @@ check a t = case (a,t) of
 
 -- Check a list of declarations
 checkDecls :: Decls -> Typing ()
-checkDecls (MutualDecls []) = return ()
-checkDecls (MutualDecls d)  = do
+checkDecls (MutualDecls _ []) = return ()
+checkDecls (MutualDecls l d)  = do
   a <- asks env
   let (idents,tele,ters) = (declIdents d,declTele d,declTers d)
   ind <- asks indent
   trace (replicate ind ' ' ++ "Checking: " ++ unwords idents)
   checkTele tele
-  local (addDecls (MutualDecls d)) $ do
+  local (addDecls (MutualDecls l d)) $ do
     rho <- asks env
     checks (tele,rho) ters
 checkDecls (OpaqueDecl _)      = return ()