a trick for indexed corecursor in Cubical Agda

coind.agda

{-# OPTIONS --cubical --guardedness #-}

open import Cubical.Foundations.Everything
open import Cubical.Data.Sigma.Properties

-- We are given an indexed container.
module _
  (Ix : Type)
  (A : Ix → Type)
  (B : (i : Ix) → A i → Type)
  (r : (i : Ix) (a : A i) → B i a → Ix) where

-- We denote the corresponding indexed M-type by M.
  record M (i : Ix) : Type where
    coinductive
    field
      head : A i
      tail : ∀ (b : B i head) → M (r i head b)

-- This is the data of a coalgebra. Instead of X being given as a type family over I,
-- here it is given as a type with a map to I. This is what causes problems later,
-- essentially because pt is not a definitional equality.
  module _
    (X : Type)
    (p : X → Ix)
    (h : (x : X) → A (p x))
    (t : (x : X) (b : B (p x) (h x)) → X)
    (pt : (x : X) (b : B (p x) (h x)) →
    p (t x b) ≡ r (p x) (h x) b) where

-- This naive attempt at defining the corecursor fails termination checking, because Agda
-- thinks the outer subst application looks problematic.
--    corec : (x : X) → M (p x)
--    M.head (corec x) = h x
--    M.tail (corec x) b = subst M (pt x b) (corec (t x b))

-- This reformulation of the corecursor does pass termination checking. Note that now the
-- target type M i is indexed by a free variable i, rather than by an expression p x as before.
    corec' : (x : X) (i : Ix) (e : p x ≡ i) → M i
    M.head (corec' x i e) = subst A e (h x)
    M.tail (corec' x i e) b =
      corec' (t x (fst be')) _ (snd be')
      where be' : Σ (B (p x) (h x)) λ b' → p (t x b') ≡ r i (subst A e (h x)) b
      -- the proof below is just path induction, but it becomes verbose in Cubical Agda
            be' =
              J
                (λ i' e' → ∀ b₀ →
                  Σ (B (p x) (h x)) λ b' → p (t x b') ≡ r i' (subst A e' (h x)) b₀)
                (subst (λ a → ∀ (b₀ : B (p x) a) → Σ (B (p x) (h x)) (λ b' → p (t x b') ≡ r (p x) a b₀))
                  (sym (substRefl {B = A} (h x)))
                  λ b₀ → b₀ , pt x b₀)
                e b

-- We can also give the usual corecursor.
    corec : (x : X) → M (p x)
    corec x = corec' x _ refl