Zippers for updating matrix factorizations

Linked SO comment on zippers for editing trees:

Oftentimes there's a tight correspondence between Monad structure and this kind of tree grafting. Here's an example


data Tree a = Leaf a | Branch (Tree a) (Tree a) deriving Functor

instance Monad Tree where
  return = Leaf
  Leaf a >>= f = f a
  Branch l r >>= f = Branch (l >>= f) (r >>= f)

Where (>>=) is doing nothing more than leaf extension (tree grafting) based on some function f :: a -> Tree a.

Then we can easily conduct the grafting you're looking for

graftRight :: Eq a => a -> a -> Tree a -> Tree a
graftRight a new t = t >>= go where
  go a' | a' == a   = Node a new
        | otherwise = Leaf a'

But that's terrifically inefficient since it'll visit every Leaf in your tree searching for the specific one you'd like to replace. We can do better if we know more information. If the tree is somehow ordered and sorted then you can use fingertree or splaytree to conduct an efficient replacement. If we know the node we'd like to replace by its path alone we can use a Zipper.

data TreeDir = L | R
data ZTree a = Root 
             | Step TreeDir (Tree a) (ZTree a)

Which lets us step in and out of the root of a Tree


stepIn :: Tree a -> (Tree a, ZTree a)
stepIn t = (t, Root)

stepOut :: (Tree a, ZTree a) -> Maybe (Tree a)
stepOut (t, Root) = Just t
stepOut _         = Nothing

and once we're inside, walk in whatever direction we like


left :: (Tree a, ZTree a) -> Maybe (Tree a, ZTree a)
left (Leaf a, zip) = Nothing
left (Branch l r, zip) = Just (l, Step R r zip)

right :: (Tree a, ZTree a) -> Maybe (Tree a, ZTree a)
right (Leaf a, zip) = Nothing
right (Branch l r, zip) = Just (r, Step L l zip)

up :: (Tree a, ZTree a) -> Maybe (Tree a, ZTree a)
up (tree, Root) = Nothing
up (tree, Step L l zip) = Just (branch l tree, zip)
up (tree, Step R r zip) = Just (branch tree r, zip)

And edit leaves

graft :: (a -> Tree a) -> (Tree a, ZTree a) -> Maybe (Tree a, ZTree a)
graft f (Leaf a, zip) = Just (f a, zip)
graft _ _             = Nothing

Or perhaps all of the leaves below a certain location using our bind from above!

graftAll :: (a -> Tree a) -> (Tree a, ZTree a) -> (Tree a, ZTree a)
graftAll f (tree, zip) = (tree >>= f, zip)

And then we can walk down to any point in the tree before doing our graft

graftBelow :: (a -> Tree a) -> [TreeDir] -> Tree a -> Maybe (Tree a)
graftBelow f steps t = perform (stepIn t) >>= stepOut where
  perform =     foldr (>=>) Just (map stepOf steps)          -- walk all the way down the path
            >=> (Just . graftAll f)                      -- graft here
            >=> foldr (>=>) Just (map (const up) steps)      -- walk back up it
  stepOf L = left
  stepOf R = right

>>> let z = Branch (Branch (Leaf "hello") (Leaf "goodbye"))
                   (Branch (Branch (Leaf "burrito")
                                   (Leaf "falcon"))
                           (Branch (Leaf "taco")
                                   (Leaf "pigeon")))

>>> graftBelow Just [] z == z
True

>>> let dup a = Branch (Leaf a) (Leaf a)
>>> graftBelow dup [L, R] z
Just (Branch (Branch (Leaf "hello") 
                     (Branch (Leaf "goodbye") 
                             (Leaf "goodbye"))) 
             (Branch (Branch (Leaf "burrito") (Leaf "falcon")) 
                     (Branch (Leaf "taco") (Leaf "pigeon"))))

>>> graftBelow dup [L, R, R] z
Nothing

ocramz / sparse-linear-algebra

Zippers for updating matrix factorizations #55