How can you build a coinductive memoization table for recursive functions over binary trees?

Question

The StreamMemo library for Coq illustrates how to memoize a function f : nat -> A over the natural numbers. In particular when f (S n) = g (f n), the imemo_make shares the computation of recursive calls.

Suppose instead of natural numbers, we want to memoize recursive functions over binary trees:

Inductive binTree : Set := | Leaf : binTree | Branch : binTree -> binTree -> binTree.

Suppose we have a function f : binTree -> A that is structurally recursive, meaning that there is a function g : A -> A -> A such that f (Branch x y) = g (f x) (f y). How do we build a similar memo table for f in Coq such that the recursive computations are shared?

In Haskell, it is not too hard to build such a memo table (see MemoTrie for example) and tie-the-knot. Clearly such memo tables are productive. How can we arrange things to convince a dependently typed language to accept such knot tying is productive?

Although I've specified the problem in Coq, I wouldd be happy with an answer in Agda or any other dependently typed language as well.

Answer 1

It's easy enough to get the recursion pattern to work with sized types. Hopefully the sharing is preserved through compilation![1]

module _ where

open import Size
open import Data.Nat

data BT (i : Size) : Set where
  Leaf : BT i
  Branch : ∀ {j : Size< i} → BT j → BT j → BT i

record Memo (A : Set) (i : Size) : Set where
  coinductive
  field
    leaf : A
    branch : ∀ {j : Size< i} → Memo (Memo A j) j

open Memo

trie : ∀ {i} {A} → (BT i → A) → Memo A i
trie f .leaf = f Leaf
trie f .branch = trie (\ l → trie \ r → f (Branch l r))

untrie : ∀ {i A} → Memo A i → BT i → A
untrie m Leaf         = m .leaf
untrie m (Branch l r) = untrie (untrie (m .branch) l) r

memo : ∀ {i A} → (BT i → A) → BT i → A
memo f = untrie (trie f)

memoFix : ∀ {A : Set} → A → (A → A → A) → ∀ {i} → BT i → A
memoFix {A} lf br = go
 where
  go h : ∀ {i} → BT i → A
  go = memo h
  h Leaf = lf
  h (Branch l r) = br (go l) (go r)

[1] https://github.com/agda/agda/issues/2918

Answer 2

I have created a "solution" that recursively memoizes structurally recursive functions of binary trees in Coq. My gist is at https://gist.github.com/roconnor/286d0f21af36c2e97e74338f10a4315b.

It operates similarly to Saizan's solution, by stratifying binary trees based on a size metric, in my case counting the number of internal nodes of binary trees, and producing a lazy stream of containers of all solutions for a particular size. Sharing happens because of a let statement in the stream generator that holds the initial part of the stream to be used in later parts of the stream.

Examples show that for vm_compute, evaluating a perfect binary tree with 8 levels after having evaluated a perfect binary tree with 9 levels is much faster than only evaluating a perfect binary tree with 8 levels.

However, I'm hesitant to accept this answer because the overhead of this particular solution is bad that it performs much worse than my memoization without structural recursion for my examples of practical inputs. Naturally, I want a solution that performs better under reasonable inputs.

I have some further comments at "[Coq-Club] How can you build a coinductive memoization table for recursive functions over binary trees?".