Some fine tuning of instance indexing

src/full/Agda/TypeChecking/DiscrimTree.hs

@@ -105,10 +105,15 @@ tickExplore tm = whenProfile Profile.Instances do
     Level{}    -> tick "explore: Level"
     MetaV{}    -> tick "explore: Meta"
     DontCare{} -> tick "explore: DontCare"
-    Dummy s _
-      | s == "eta-record argument in instance head" -> tick "explore: from eta-expansion"
     _ -> pure ()
 
+-- | Checks whether the given abstraction could have been 'NoAbs'.
+isNoAbs :: Abs Term -> Maybe Term
+isNoAbs (NoAbs _ b) = Just b
+isNoAbs (Abs _ b)
+  | not (0 `freeIn` b) = Just (strengthen __IMPOSSIBLE__ b)
+  | otherwise          = Nothing
+
 -- | Split a term into a 'Key' and some arguments. The 'Key' indicates
 -- whether or not the 'Term' is in head-normal form, and provides a
 -- quick way to match on the head.
@@ -119,7 +124,7 @@ tickExplore tm = whenProfile Profile.Instances do
 -- Presently, non-head-normal terms end up with an empty argument list.
 splitTermKey :: Bool -> Int -> Term -> TCM (Key, [Term])
 splitTermKey precise local tm = fmap (fromMaybe (FlexK, [])) . runMaybeT $ do
-  (b, tm') <- ifBlocked tm (\_ _ -> mzero) (curry pure)
+  (b, tm') <- ifBlocked tm (\_ _ -> mzero) (\b -> fmap (b,) . constructorForm)
 
   case tm' of
     Def q as | ReallyNotBlocked <- b -> do
@@ -143,18 +148,29 @@ splitTermKey precise local tm = fmap (fromMaybe (FlexK, [])) . runMaybeT $ do
       (arity, as) <- termKeyElims precise ty as
       pure (RigidK q arity, as)
 
-    Pi dom ret
-      -- For slightly more accurate matching, we decompose non-dependent
-      -- 'Pi's into a distinguished key.
-      | NoAbs _ b <- ret -> do
-        whenProfile Profile.Instances $ tick "funk: non-dependent function"
-        pure (FunK, [unEl (unDom dom), unEl b])
-
-      | otherwise -> do
-        whenProfile Profile.Instances $ tick "funk: genuine pi"
-        pure (PiK, [])
-
-    _ -> pure (FlexK, [])
+    Pi dom ret ->
+      let
+        -- If we're looking at a non-dependent function type, then we
+        -- might as well represent the codomain accurately; Otherwise,
+        -- turn the codomain into a wildcard.
+        ret' = case isNoAbs (unEl <$> ret) of
+          Just b  -> b
+          Nothing -> __DUMMY_TERM__
+      in pure (PiK, [unEl (unDom dom), ret'])
+
+    Lam _ body
+      -- Constant lambdas come up quite a bit, particularly (in cubical
+      -- mode) as the domain of a PathP. Having this trick improves the
+      -- indexing of 'Dec' instances in the 1Lab significantly.
+      | Just b <- isNoAbs body -> pure (ConstK, [b])
+
+    -- Probably not a good idea for accurate indexing if universes
+    -- overlap literally everything else.
+    Sort _ -> pure (SortK, [])
+
+    _ -> do
+      reportSDoc "tc.instance.split" 30 $ pretty tm
+      pure (FlexK, [])
 
 termPath :: Int -> [Key] -> [Term] -> TCM [Key]
 termPath local acc []        = pure $! reverse acc
@@ -188,12 +204,15 @@ insertDT local key val tree = ignoreAbstractMode do
     ]
   pure $ mergeDT it tree
 
+-- | If a term matches this key, how many arguments does it place on the
+-- spine?
 keyArity :: Key -> Int
 keyArity = \case
   RigidK _ a -> a
   LocalK _ a -> a
-  FunK       -> 2
-  PiK        -> 0
+  PiK        -> 2
+  ConstK     -> 1
+  SortK      -> 0
   FlexK      -> 0
 
 -- | Look up a 'Term' in the given discrimination tree. The returned set
@@ -297,21 +316,13 @@ lookupDT term tree = ignoreAbstractMode (match [term] tree) where
         -- Actually take the branch corresponding to our rigid head.
         branch <- visit k sp'
 
-        -- Function values get unpacked to their components on the
-        -- spine, but proper Π types don't. Since Π-type instances also
-        -- apply to function types, we have to explore that branch too,
-        -- if our rigid head is a function.
-        funIsPi <- case k of
-          FunK -> visit PiK (sp0 ++ sp1)
-          _    -> pure mempty
-
         -- When exploring the rest of the tree, the value we cased on
         -- has to be put back in the tree. mergeDT does not perform
         -- commuting conversions to ensure that variables aren't
         -- repeatedly cased on.
         rest <- match ts rest
 
-        pure $! rest <> branch <> funIsPi
+        pure $! rest <> branch
 
   match ts tree@(CaseDT i _ rest) = do
     reportSDoc "tc.instance.discrim.lookup" 99 $ vcat

src/full/Agda/TypeChecking/DiscrimTree/Types.hs

@@ -22,11 +22,18 @@ data Key
     -- postulates) identified by a QName.
   | LocalK {-# UNPACK #-} !Int {-# UNPACK #-} !Int
     -- ^ Local variables.
-  | FunK
-    -- ^ Non-dependent function types are represented accurately.
+
   | PiK
-    -- ^ Dependent function types get their own branch in the tree, but
-    -- their arguments are not considered.
+    -- ^ Dependent function types. The domain will be represented
+    -- accurately, for the case of a genuine dependent function type,
+    -- the codomain will be a dummy.
+
+  | ConstK
+    -- ^ Constant lambdas.
+
+  | SortK
+    -- ^ Universes.
+
   | FlexK
   -- ^ Anything else.
   deriving (Show, Eq, Ord, Generic)

src/full/Agda/TypeChecking/Monad/Statistics.hs

@@ -13,9 +13,13 @@ import Control.Monad.State
 import Control.Monad.Writer
 import Control.Monad.Trans.Maybe
 
+import Data.IORef
+
 import qualified Data.Map as Map
 import qualified Text.PrettyPrint.Boxes as Boxes
 
+import System.IO.Unsafe
+
 import Agda.Syntax.TopLevelModuleName (TopLevelModuleName)
 
 import Agda.TypeChecking.Monad.Base
@@ -58,13 +62,17 @@ instance MonadStatistics TCM where
       update  = Map.insertWith (\ new old -> f old) x dummy
       dummy   = f 0
 
+statistics :: IORef Statistics
+statistics = unsafePerformIO (newIORef mempty)
+{-# NOINLINE statistics #-}
+
 -- | Get the statistics.
-getStatistics :: ReadTCState m => m Statistics
-getStatistics = useR stStatistics
+getStatistics :: TCM Statistics
+getStatistics = liftIO $ readIORef statistics
 
 -- | Modify the statistics via given function.
 modifyStatistics :: (Statistics -> Statistics) -> TCM ()
-modifyStatistics f = stStatistics `modifyTCLens` f
+modifyStatistics f = liftIO $ modifyIORef' statistics f
 
 -- | Increase specified counter by @1@.
 tick :: MonadStatistics m => String -> m ()

src/full/Agda/TypeChecking/Pretty.hs

@@ -587,8 +587,9 @@ instance PrettyTCM Key where
   prettyTCM = \case
     RigidK q a -> prettyTCM q <> superscript a
     LocalK i a -> "@" <> pretty i <> superscript a
-    FunK       -> "Fun"
     PiK        -> "Pi"
+    ConstK     -> "Const"
+    SortK      -> "Sort"
     FlexK      -> "_"
 
 {-# SPECIALIZE prettyTCM :: Key -> TCM Doc #-}

src/full/Agda/TypeChecking/Serialise.hs

@@ -77,7 +77,7 @@ import Agda.Utils.Impossible
 -- 32-bit machines). Word64 does not have these problems.
 
 currentInterfaceVersion :: Word64
-currentInterfaceVersion = 20240222 * 10 + 2
+currentInterfaceVersion = 20240225 * 10 + 0
 
 -- | The result of 'encode' and 'encodeInterface'.
 

src/full/Agda/TypeChecking/Serialise/Instances/Internal.hs

@@ -659,16 +659,18 @@ instance EmbPrj RemoteMetaVariable where
 instance EmbPrj Key where
   icod_ (RigidK x y) = icodeN 0 RigidK x y
   icod_ (LocalK x y) = icodeN 1 LocalK x y
-  icod_ FunK         = icodeN 2 FunK
-  icod_ PiK          = icodeN 3 PiK
-  icod_ FlexK        = icodeN 4 FlexK
+  icod_ PiK          = icodeN 2 PiK
+  icod_ FlexK        = icodeN 3 FlexK
+  icod_ ConstK       = icodeN 4 ConstK
+  icod_ SortK        = icodeN 5 SortK
 
   value = vcase valu where
     valu [0, x, y] = valuN RigidK x y
     valu [1, x, y] = valuN LocalK x y
-    valu [2]       = valuN FunK
-    valu [3]       = valuN PiK
-    valu [4]       = valuN FlexK
+    valu [2]       = valuN PiK
+    valu [3]       = valuN FlexK
+    valu [4]       = valuN ConstK
+    valu [5]       = valuN SortK
     valu _         = malformed
 
 instance (EmbPrj a, Ord a) => EmbPrj (DiscrimTree a) where