A space leak due to absName

[nad]

I noticed that use of absName can lead to space leaks. Biographical and retainer profiling (for a certain kind of contrived test case) indicated that telViewUpTo' retained lots of data that was never used:

agda/src/full/Agda/TypeChecking/Telescope.hs

Lines 374 to 381 in f787254

	telViewUpTo' :: (MonadReduce m, MonadAddContext m) => Int -> (Dom Type -> Bool) -> Type -> m TelView
	telViewUpTo' 0 p t = return $ TelV EmptyTel t
	telViewUpTo' n p t = do
	t <- reduce t
	case unEl t of
	Pi a b | p a -> absV a (absName b) <$> do
	underAbstractionAbs a b $ \b -> telViewUpTo' (n - 1) p b
	_ -> return $ TelV EmptyTel t

In the penultimate case a and b are used after some potentially expensive computation. What does absV do?

agda/src/full/Agda/TypeChecking/Substitute.hs

Lines 1217 to 1218 in f787254

	absV :: Dom a -> ArgName -> TelV a -> TelV a
	absV a x (TelV tel t) = TelV (ExtendTel a (Abs x tel)) t

The argument a is returned, but not all of b, only absName b. I guessed that all of b is retained, and inspection of the Core for telViewUpTo' (for GHC 9.4.2 without -foptimise-heavily) suggests that this is the case:

                         fmap
                           $dFunctor1
                           (let {
                              x = case b of {
                                    Abs ds5 ds6 -> ds5;
                                    NoAbs ds5 ds6 -> ds5
                                  } } in
                            \ ds5 ->
                              case ds5 of { TelV tel t2 -> TelV (ExtendTel a (Abs x tel)) t2 })
                           (underAbstractionAbs $dSubst1 $dMonadAddContext a b lvl121)

(A specialised variant of the code is similar.) If we replace the code above with the following one, then it appears as if only the name is retained:

    Pi a b | p a -> let !bn = absName b in
                    absV a bn <$> do
                      underAbstractionAbs a b $ \b -> telViewUpTo' (n - 1) p b

                         join {
                           $j bn
                             = case bn of bn1 { __DEFAULT ->
                               fmap
                                 $dFunctor1
                                 (\ ds5 ->
                                    case ds5 of { TelV tel t2 ->
                                    TelV (ExtendTel a (Abs bn1 tel)) t2
                                    })
                                 (underAbstractionAbs $dSubst1 $dMonadAddContext a b lvl121)
                               } } in
                         case b of {
                           Abs ds5 ds6 -> jump $j ds5;
                           NoAbs ds5 ds6 -> jump $j ds5
                         }

This change led to a speed-up of more than 50% for my test case, even though the slower variant of Agda was compiled using -foptimise-heavily, and the faster one wasn't.

Here is the implementation of Abs:

agda/src/full/Agda/Syntax/Internal.hs

Lines 241 to 245 in f787254

	data Abs a = Abs { absName :: ArgName, unAbs :: a }
	-- ^ The body has (at least) one free variable.
	-- Danger: 'unAbs' doesn't shift variables properly
	| NoAbs { absName :: ArgName, unAbs :: a }
	deriving (Data, Functor, Foldable, Traversable, Generic)

Note that absName is a field of two different constructors. I replaced this definition by the following one, and got a similar speed-up without forcing the name (after fixing one line of code that used the pattern Abs{unAbs=ts}):

data Abs a = AbsOrNoAbs
  { absKind :: !Bool
  , absName :: ArgName
  , unAbs   :: a
  }
  deriving (Data, Functor, Foldable, Traversable, Generic)

pattern Abs, NoAbs :: ArgName -> a -> Abs a
pattern Abs   n x = AbsOrNoAbs True  n x
pattern NoAbs n x = AbsOrNoAbs False n x

{-# COMPLETE Abs, NoAbs #-}

However, note that in this case the Core code still suggests that b is retained:

                         fmap
                           $dFunctor1
                           (let { x = case b of { AbsOrNoAbs ds5 ds6 ds7 -> ds6 } } in
                            \ ds5 ->
                              case ds5 of { TelV tel t2 ->
                              TelV (ExtendTel a (AbsOrNoAbs True x tel)) t2
                              })
                           (underAbstractionAbs $dSubst1 $dMonadAddContext a b lvl123)

Here is an excerpt of the specialised variant of the code:

                     case ((((underAbstractionAbs
                                $dSubst1 $fMonadAddContextTCMT a b lvl122)
                             `cast` <Co:4> :: TCMT IO (TelV Type)
                                              ~R# (IORef TCState -> TCEnv -> IO (TelV Type)))
                              (wild `cast` <Co:3> :: STRef RealWorld TCState ~R# IORef TCState)
                              e)
                           `cast` <Co:3> :: IO (TelV Type)
                                            ~R# (State# RealWorld
                                                 -> (# State# RealWorld, TelV Type #)))
                            ipv2
                     of
                     { (# ipv4, ipv5 #) ->
                     (# ipv4,
                        case ipv5 of { TelV tel t1 ->
                        TelV
                          (ExtendTel
                             a
                             (AbsOrNoAbs
                                True (case b of { AbsOrNoAbs ds6 ds7 ds8 -> ds7 }) tel))
                          t1
                        } #)
                     }

I'm not sure exactly what is going on. Is there some later pass of GHC that optimises this code further?

What do you think? Should we use the bang pattern, or the new definition of Abs? The new definition of Abs might be helpful in other places, but I'm not sure exactly why it works.

(Before we commit to any change I think it makes sense to also run the standard library benchmark.)

Edit: I renamed Abs' to AbsOrNoAbs (The Haskell code was not consistent with the Core code.)

[andreasabel]

How about making absName a strict field? It will still be a pointer to a string then, so sharing shouldn't be lost. (Correct me if I am wrong, I am not super firm on these issues.)
(Even in Abs', it could be a strict field, I suppose.)

[nad]

How about making absName a strict field?

I don't see why that would help, and a quick test suggests that it doesn't.

However, I also don't see why the switch from two constructors to one helps. If the type was unlifted, then it would be safe to perform the projection eagerly, but the type is lifted, so values of that type can be thunks (with one or two constructors). I tried to use -fpedantic-bottoms, and still got the speed-up. Can anyone explain what is going on?

[nad]

(Before we commit to any change I think it makes sense to also run the standard library benchmark.)

I type-checked Data.Nat (with Agda compiled without -foptimise-heavily) using both approaches, as well as without any change. When I changed Abs Agda seemed to become a little slower. Forcing the name seemed to work better, so I've decided to go with that approach.