Typing TABA

Just a short one again today!

There’s an excellent talk by Kenneth Foner at Compose from 2016 which goes through a paper by Danvy and Goldberg (2005) called “There and Back Again” (or TABA). You should watch the talk and read the paper if you’re in any way excited by the weird and wonderful algorithms we use in functional languages to do simple things like reversing a list.

The function focused on in the paper is one which does the following:

zipRev :: [a] -> [b] -> [(a,b)]
zipRev xs ys = zip xs (reverse ys)

But does it in one pass, without reversing the second list. It uses a not-insignificant bit of cleverness to do it, but you can actually arrive at the same solution in a pretty straightforward way by aggressively converting everything you can to a fold. The result is the following:

zipRev :: [a] -> [b] -> [(a,b)]
zipRev xs ys = foldl f b ys xs
  where
    b _ = []
    f k y (x:xs) = (x,y) : k xs

I have written a little more on this function and the general technique before.

The talk goes through the same stuff, but takes a turn then to proving the function total: our version above won’t work correctly if the lists don’t have the same length, so it would be nice to provide that guarantee in the types somehow. Directly translating the version from the TABA paper into one which uses length-indexed vectors will require some nasty, expensive proofs, though, which end up making the whole function quadratic. The solution in the talk is to call out to an external solver which gives some extremely slick proofs (and a very nice interface). However, yesterday I realised you needn’t use a solver at all: you can type the Haskell version just fine, and you don’t even need the fanciest of type-level features.

As ever, the solution is another fold. To demonstrate this rather short solution, we’ll first need the regular toolbox of types:

data Nat = Z | S Nat

data Vec (a :: Type) (n :: Nat) where
  Nil :: Vec a Z
  (:-) :: a -> Vec a n -> Vec a (S n)

And now we will write a length-indexed left fold on this vector. The key trick here is that the type passed in the recursive call changes, by composition:

newtype (:.:) (f :: b -> Type) (g :: a -> b) (x :: a) = Comp { unComp :: f (g x) }

Safe coercions will let us use the above type safely without a performance hit, resulting in the following linear-time function:

foldlVec :: forall a b n. (forall m. a -> b m -> b (S m)) -> b Z -> Vec a n -> b n
foldlVec f b Nil = b
foldlVec f b (x :- xs) = unComp (foldlVec (c f) (Comp (f x b)) xs)
  where
    c :: (a -> b (S m) -> b (S (S m))) -> (a -> (b :.: S) m -> (b :.: S) (S m))
    c = coerce
    {-# INLINE c #-}

We can use this function to write vector reverse:

reverseVec :: Vec a n -> Vec a n
reverseVec = foldlVec (:-) Nil

Now, to write the reversing zip, we need another newtype to put the parameter in the right place, but it is straightforward other than that.

newtype VecCont a b n = VecCont { runVecCont :: Vec a n -> Vec (a,b) n }

revZip :: Vec a n -> Vec b n -> Vec (a,b) n
revZip = flip $ runVecCont . 
  foldlVec
      (\y k -> VecCont (\(x :- xs) -> (x,y) :- runVecCont k xs))
      (VecCont (const Nil))