The trek goes on

0025

Termination as type casting

Posted at 22:44 on 20 December 2021, and revised at 17:05 on 27 January 2022

There is a reduction in blogging frequency because I have been working with Liang-Ting and Viktor on a datatype-generic library that specialises to native Agda programs using reflection. We aim to submit a paper about this to ICFP ’22 (whose deadline is somewhat earlier than expected), and we’re going as fast as we can without getting overstretched. I’m taking a small break by writing this post (which is therefore a bit rushed) about an Agda exercise I finished a while ago.


I mentioned in 0023 that Zong-You and I have been playing with tree traversals; it began with Conor McBride’s ‘clowns and jokers’ paper, which presented a stack machine for computing folds (which I don’t suppose was new, but the technique for deriving the type of stack elements using dissection was). Later we found out that Carlos Tomé Cortiñas and Wouter Swierstra did a complete Agda formalisation of the construction, in particular making the termination of machine execution evident to Agda. The formalisation looks unsatisfactory to me for two reasons: (i) it’s still a function that goes directly from input to output rather than allowing suspension and resumption in the middle of the computation, which is one key benefit of using the stack machine instead of the original fold function, and (ii) termination is made evident by heavily altering and complicating the definition, which looks more like a necessary evil (due to Agda’s insistence on totality).

To allow suspension and resumption, we should model the stack machine computation as a state transition function, which can then be invoked as many or few times as one wants, stopping at various points of interest during the computation. More specifically, the user gets and traverses a colist of states computed iteratively by the state transition function. This colist can be truly infinite in general: an observation here is that the stack machine works by gradually decomposing the input tree, and decomposition is the coinductive view, so the stack machine in fact works on codata, which can be infinite, in which case the computation is infinite too. This also explains why termination wasn’t evident for the original definition, simply because the original definition is in general non-terminating. On the other hand, if the input tree is finite, then the computation should terminate, which, for the colist formulation, means that we can cast the potentially infinite colist into a finite list, from which the last state containing the final result can be drawn. The rest of this post formalises these ideas.

First, a brief recap on the fold machine. Suppose that we are folding internally labelled binary trees, whose shape is defined by the base functor

data ITreeF (A X : Set) : Set where
  leaf : ITreeF A X
  node : A → X → X → ITreeF A X

of which we can take the least and greatest fixed points,

data μITree (A : Set) : Set where
  con : ITreeF A (μITree A) → μITree A

record νITree (A : Set) : Set where
  coinductive
  field decon : ITreeF A (νITree A)

open νITree

and we can convert the least fixed point to the greatest fixed point,

toνITree : μITree A → νITree A
decon (toνITree (con  leaf       )) = leaf
decon (toνITree (con (node a t u))) = node a (toνITree t) (toνITree u)

although we’ll only be working on the greatest fixed point for now. The machine state type is parametrised by the tree element/label type A and a result type B, and each state consists of a focus and a stack (whose type is just a renamed version of the standard library list type):

record MachineState (A B : Set) : Set where
  constructor ⟨_,_⟩
  field
    focus : Focus A B
    stack : Stack (∂ITree A B)

The content of the focus depends on the direction in which the machine is travelling: if the machine is travelling downwards, then it is decomposing the input tree, in which case the focus is on a sub-tree to be decomposed; if the machine is travelling upwards, then it has finished processing a sub-tree, in which case the focus stores the sub-result computed from the sub-tree.

data Focus (A B : Set) : Set where
  ↓_ : νITree A → Focus A B
  ↑_ : B        → Focus A B

The stack stores contents that are out of focus: When the machine focuses on a sub-tree, the right sub-tree is stored in the stack so that it can be popped out and traversed after the sub-tree in focus is processed. Since we’re computing a fold, an algebra of type A → B → B → B is used at each node to combine the results computed from the sub-trees into the result for the node. I choose to partially apply the algebra as soon as the machine sees a node and its label or the result computed from the left sub-tree, and store the rest of the algebra in the stack too.

data ∂ITree (A B : Set) : Set where
  ○-⟨_,_⟩ : νITree A → (B → B → B) → ∂ITree A B
  ⟨_⟩-○   :                (B → B) → ∂ITree A B

To traverse a tree, we prepare an initial state where the focus contains the tree and the stack is empty. Then we run the machine by repeatedly invoking the state transition function:

step : B → (A → B → B → B) → MachineState A B → Maybe (MachineState A B)
step b f ⟨ ↓ t , stk ⟩ with νITree.decon t
... | leaf       = just ⟨ ↑ b ,                 stk ⟩
... | node a t u = just ⟨ ↓ t , ○-⟨ u , f a ⟩ ∷ stk ⟩
step _ _ ⟨ ↑ b , [] ⟩ = nothing
step _ _ ⟨ ↑ b , ○-⟨ u , f₁ ⟩ ∷ stk ⟩ = just ⟨ ↓ u    , ⟨ f₁ b ⟩-○ ∷ stk ⟩
step _ _ ⟨ ↑ b , ⟨ f₂ ⟩-○     ∷ stk ⟩ = just ⟨ ↑ f₂ b ,              stk ⟩

When travelling downwards, if the focussed tree is a leaf, we return the base value b : B upwards; if the tree is a node, we go downwards into the left sub-tree and put the right sub-tree along with the algebra f partially applied to the node label in the stack for later processing. When travelling upwards, if there are nothing else to process in the stack, the machine halts; otherwise we pop out a sub-tree (if any) to continue processing and a partially applied algebra to apply to the result in focus.

To run the machine, we unfold a colist of states with step. Again start with the base functor for lists

data ListF (A X : Set) : Set where
  []  : ListF A X
  _∷_ : A → X → ListF A X

and its least and greatest fixed points

data μList (A : Set) : Set where
  con : ListF A (μList A) → μList A

record νList (A : Set) : Set where
  coinductive
  field decon : ListF A (νList A)

open νList

also with conversion from the former to the latter:

toνList : μList A → νList A
decon (toνList (con []      )) = []
decon (toνList (con (x ∷ xs))) = x ∷ toνList xs

We can define the standard unfold operator on colists,

unfoldr : (B → ListF A B) → B → νList A
decon (unfoldr f b) with f b
... | []     = []
... | a ∷ b' = a ∷ unfoldr f b'

using which we can then define a function that can be used to apply step iteratively to an initial state and unfold a colist of states:

iterateMaybe : (A → Maybe A) → Maybe A → νList A
iterateMaybe f = unfoldr (maybe (λ a → a ∷ f a) [])

(Incidentally, this seems to be a less known but more natural and general version of iterateMaybe than the version of type (A → Maybe A) → A → νList A, which can only produce non-empty colists.) Now the function

run : B → (A → B → B → B) → νITree A → νList (MachineState A B)
run b f t = iterateMaybe (step b f) (just ⟨ ↓ t , [] ⟩)

gives the colist of states produced by the machine.

Our goal is to show that the νList above is in fact finite by casting it to a μList. While it is always possible to cast a μList to a νList, as witnessed by toνList, the opposite direction requires some condition. The condition I come up with is that the νList being cast should be simulated by a finite νList of the form toνList xs where xs is some μList. Let A ~ B = A → B → Set be the type of binary relations on A and B : Set. When I say that a colist as : νList A is simulated by another colist bs : νList A’ with respect to some relation R : A ~ B, I mean that whenever as deconstructs to a head a and a tail as', bs also deconstructs to a head b and a tail bs' such that R a b holds and as' is again simulated by bs' with respect to R. Intuitively, if as is simulated by bs, then the length of as is bounded by the length of bs (and a binary relation holds for their corresponding elements), so if we know bs is finite, then as has to be finite as well, and we can successfully cast as to a μList. Writing all these down formally:

data ListSimF (R : A ~ B) (S : X ~ Y) : ListF A X ~ ListF B Y where
  []  : {bs : ListF B Y} → ListSimF R S [] bs
  _∷_ : {a : A} {b : B} → R a b →
        {x : X} {y : Y} → S x y → ListSimF R S (a ∷ x) (b ∷ y)

record νListSim (R : A ~ B) (as : νList A) (bs : νList B) : Set where
  coinductive
  field decon : ListSimF R (νListSim R) (decon as) (decon bs)

open νListSim

toμList : {R : A ~ B} (as : νList A) (bs : μList B)
        → νListSim R as (toνList bs) → μList A
toμList as _           sim with decon as | decon sim
toμList _  _             _ | []     | _       = con []
toμList _ (con (_ ∷ bs)) _ | a ∷ as | _ ∷ sim = con (a ∷ toμList as bs sim)

Which list are we going to use to bound the length of run b f (toνITree t) given some t : μITree A? The simplest possible: a list of units whose length is the number of steps we know the machine will take before it halts. But we can only count the steps if we know that all the trees in the machine states have finite sizes (or otherwise the machine doesn’t terminate). Our strategy is to construct an abstract version of the machine where all the trees in the states are finite, and show that it simulates the original machine when running on a finite tree; we will also show that the colist of states produced by the abstract version of the machine is simulated by the aforementioned list of units, so by transitivity the colist of states produced by the original machine is also simulated by the list of units, and then we will have shown that the original machine terminates on a finite tree.

The abstract machine is structurally the same as the original one but missing some contents since we don’t care about computing the results:

record MachineState' (A : Set) : Set where
  constructor ⟨_,_⟩'
  field
    focus : Maybe (μITree A)
    stack : Stack (Maybe (μITree A))

step' : MachineState' A → Maybe (MachineState' A)
step' ⟨ just (con  leaf       ) , stk ⟩' = just ⟨ nothing ,          stk ⟩'
step' ⟨ just (con (node a t u)) , stk ⟩' = just ⟨ just t ,  just u ∷ stk ⟩'
step' ⟨ nothing ,                 []  ⟩' = nothing
step' ⟨ nothing ,       just u  ∷ stk ⟩' = just ⟨ just u , nothing ∷ stk ⟩'
step' ⟨ nothing ,       nothing ∷ stk ⟩' = just ⟨ nothing ,          stk ⟩'

All the trees in a MachineState' are finite, and we can calculate the precise number of steps that the machine will take before it halts with these functions:

countTreeSteps : μITree A → ℕ → ℕ
countTreeSteps (con  leaf       ) = suc
countTreeSteps (con (node a t u)) =
  suc ∘ countTreeSteps t ∘ suc ∘ countTreeSteps u ∘ suc

countStackSteps : Stack (Maybe (μITree A)) → ℕ
countStackSteps            []   = zero
countStackSteps (just t  ∷ stk) =
           suc (countTreeSteps t (suc (countStackSteps stk)))
countStackSteps (nothing ∷ stk) = suc (countStackSteps stk)

countSteps : MachineState' A → ℕ
countSteps ⟨ just t  , stk ⟩' = countTreeSteps t (countStackSteps stk)
countSteps ⟨ nothing , stk ⟩' =                   countStackSteps stk

countTreeSteps calculates the number of steps (in the form of ‘difference natural numbers’ of type ℕ → ℕ to avoid associativity issues in later proofs) the machine will take to reduce a focussed tree to a result: if the tree is a leaf, then it takes one step (from focussing on the leaf (and going downwards) to focussing on the result for the leaf (and going upwards)); if it is a node, then there are

  • one step for focussing on the left sub-tree,
  • several steps for reducing the left sub-tree to a sub-result,
  • one step for switching the focus to the right sub-tree,
  • several steps for reducing the right sub-tree to a sub-result, and
  • one final step for combining the two sub-results.

Similarly, countStackSteps calculates the number of steps for eliminating the stack: for example, in the just t ∷ stk case, there are

  • one step for popping out t (which is a right sub-tree) and putting it in focus,
  • several steps for reducing t to a sub-result,
  • one step for combining the result with the sub-result for some left sub-tree, and
  • several steps for eliminating the rest of the stack.

Finally, countSteps calculates the number of steps to reach the final state from the given state by summing the numbers of steps it will take to eliminate the focus and the stack. Then it is straightforward to establish the simulation by basically instructing Agda to look into the definitions:

terminationSim :
    (ms' : MachineState' A)
  → νListSim (λ _ _ → ⊤) (iterateMaybe step' (just ms'))
                         (toνList (replicate (suc (countSteps ms')) tt))
decon (terminationSim ⟨ just (con leaf        ) , stk ⟩') =
  tt ∷ terminationSim _
decon (terminationSim ⟨ just (con (node a t u)) , stk ⟩') =
  tt ∷ terminationSim _
decon (terminationSim ⟨ nothing ,                 []  ⟩') =
  tt ∷ record { decon = [] }
decon (terminationSim ⟨ nothing ,        just u ∷ stk ⟩') =
  tt ∷ terminationSim _
decon (terminationSim ⟨ nothing ,       nothing ∷ stk ⟩') =
  tt ∷ terminationSim _

where replicate : ℕ → A → μList A produces a list of the specified length consisting of the given element.

The other simulation requires setting up an abstraction relation for relating an original state and an abstract state, saying that they’re basically the same; in particular, the potentially infinite trees in the original state are actually the necessarily finite trees in the abstract state. I’ll just dump the definitions, where the Pointwise relations from the standard library are helpful.

funR : (A → B) → (A ~ B)
funR f x y = f x ≡ y

from∂ITree : ∂ITree A B → Maybe (νITree A)
from∂ITree ○-⟨ t , _ ⟩ = just t
from∂ITree ⟨ _ ⟩-○     = nothing

fromFocus : Focus A B → Maybe (νITree A)
fromFocus (↓ t) = just t
fromFocus (↑ _) = nothing

abstractionR : MachineState A B → MachineState' A → Set
abstractionR ⟨ focus , stack ⟩ ⟨ focus' , stack' ⟩' =
  Maybe.Pointwise (funR toνITree) focus' (fromFocus focus) ×
  Stack.Pointwise (Maybe.Pointwise (funR toνITree))
                                  stack' (Stack-map from∂ITree stack)
  -- Stack-map : (A → B) → Stack A → Stack B

Then the simulation proof proceeds similarly by instructing Agda to look into the definitions; it’s kind of difficult to typeset, so I’ll just leave it in the source code (link at the end).

abstractionSim :
    (b : B) (f : A → B → B → B)
    (ms : MachineState A B) (ms' : MachineState' A) → abstractionR ms ms'
  → νListSim abstractionR (iterateMaybe (step b f) (just ms ))
                          (iterateMaybe  step'     (just ms'))

Finally, transitivity of simulation gives us the simulation that can be used to turn the colist run b f (toνITree t) of states into a finite list.

fullTerminationSim :
    (b : B) (f : A → B → B → B) (t : μITree A)
  → νListSim (λ _ _ → ⊤)
             (iterateMaybe (step b f) (just ⟨ ↓ (toνITree t) , [] ⟩))
             (toνList (replicate (suc (countTreeSteps t zero)) tt))
fullTerminationSim b f t =
  νListSim-map (λ _ → tt)
    (νListSim-trans (abstractionSim b f (⟨ ↓ (toνITree t) , [] ⟩ )
                                        (⟨ just        t  , [] ⟩')
                                        (  just      refl , []   ))
                    (terminationSim     (⟨ just        t  , [] ⟩')))

Check the source code for the full details.

Agda file: FoldMachineTermination.agda


Back to work…