A total map as a zipper of an infinite perfect binary tree

1d

 -1 0 1 2 3 4 5
\ / \\/ \ / \ /
 \   //  \   /
  \ //    \ /
   \\      /
    \\    /
     \\  /
      \\/
       //
      //

2d

|___|   |___|
| | |   | | |
  |_______|
  |   |   |
|_|_| | |_|_|
|   | | |   |
______|
      |
|___| | |___|
| | | | | | |
  |___|___|
  |       |
|_|_|   |_|_|
|   |   |   |

totalmap.hs

{-# LANGUAGE DeriveFunctor, TypeFamilies, MultiParamTypeClasses #-}

import Control.Applicative
import Control.Arrow (first)
import Control.Comonad
import Control.Comonad.Store.Class
import Data.Key
import Data.Functor.Representable


class Eq p => Pos p where
  -- cons . uncons == id 
  uncons :: p -> (p, Bool)
  cons   :: (p, Bool) -> p
  -- uncons zero == (zero, False)
  zero :: p

one :: Pos p => p
one = cons (zero, True)

conv :: (Pos p, Pos q) => p -> q
conv p | p == zero = zero
       | otherwise = cons . first conv . uncons $ p


instance Pos Integer where
  zero = 0
  uncons i = let (d, m) = i `divMod` 2 in (negate d, m == 1)
  cons (i, b) = (if b then 1 else 0) - 2 * i

instance Pos p => Pos (p, p) where
  zero = (zero, zero)
  uncons (x, y) = let (x', b) = uncons x in ((y, x'), b)
  cons ((x, y), b) = (cons (y, b), x)

instance Pos p => Pos (p, p, p) where
  zero = (zero, zero, zero)
  uncons (x, y, z) = let (x', b) = uncons x in ((y, z, x'), b)
  cons ((x, y, z), b) = (cons (z, b), x, y)

instance Pos [Bool] where
  zero = []
  uncons [] = ([], False)
  uncons (b:bs) = (bs, b)
  cons (bs, b) = b:bs


data TotalMap p a = TotalMap !p !a (TotalMapD a) deriving Functor
data TotalMapD a = TotalMapD { side :: !Bool, other :: a, ctxD :: TotalMapD (a, a) } deriving Functor

instance (Show a, Show p) => Show (TotalMap p a) where
  show (TotalMap p a _) = "TotalMap " ++ show p ++ " " ++ show a ++ " (...)"

instance Pos p => Applicative (TotalMap p) where
  pure = pureRep
  fs <*> TotalMap p a ca = let TotalMap _ f cf = seek p fs in TotalMap p (f a) $ zipWithD id cf ca
    
instance Comonad (TotalMap p) where
  extract (TotalMap _ a _) = a
  extend f z@(TotalMap p _ c) = TotalMap p (f z) $ fmap (\a -> f (TotalMap p a c)) c

instance Pos p => ComonadStore p (TotalMap p) where
  pos (TotalMap p _ _) = p
  peek p = extract . seek p
  seek tp (TotalMap sp a c) = let (a', c') = gotoD sp tp (a, c) in TotalMap tp a' c'
  seeks f w = seek (f (pos w)) w
   
type instance Key (TotalMap p) = p

instance Pos p => Lookup (TotalMap p) where lookup = lookupDefault
instance Pos p => Adjustable (TotalMap p) where
  adjust f p z = seek (pos z) . modify f . seek p $ z
instance Pos p => Indexable (TotalMap p) where
  index = flip peek
instance Pos p => Representable (TotalMap p) where
  tabulate f = TotalMap zero (f zero) (tabulateD f)


modify :: (a -> a) -> TotalMap p a -> TotalMap p a
modify m (TotalMap p a f) = TotalMap p (m a) f


gotoD :: Pos p => p -> p -> (a, TotalMapD a) -> (a, TotalMapD a)
gotoD sp tp | sp == tp = id
            | otherwise = let (tp', b) = uncons tp in down b . gotoD (fst $ uncons sp) tp' . up
  where
    up (a, TotalMapD False b c) = ((a, b), c)
    up (b, TotalMapD True  a c) = ((a, b), c)
    down False ((a, b), c) = (a, TotalMapD False b c)
    down True  ((a, b), c) = (b, TotalMapD True  a c)

tabulateD :: Pos p => (p -> a) -> TotalMapD a
tabulateD f = TotalMapD False (f one) $ tabulateD (\p -> (f (cons (p, False)), f (cons (p, True)))) 

zipWithD :: (a -> b -> c) -> TotalMapD a -> TotalMapD b -> TotalMapD c
zipWithD ap (TotalMapD _ f cf) (TotalMapD b a ca) = 
  TotalMapD b (ap f a) (zipWithD (\(f1, f2) (a1, a2) -> (ap f1 a1, ap f2 a2)) cf ca)


testp :: (Integer, Integer, Integer)
testp = (5000000000000, -500000000, 30000000000000)
test1 :: TotalMap (Integer, Integer, Integer) Integer
test1 = adjust (+ 10) (3, 3, 5) $ modify (+1) $ seek testp (tabulate (\(a, b, c) -> a + b + c))
test2 :: Integer
test2 = index test1 (3, 3, 5)