GridZipper

Here is some code I wrote to implement an idea posed by someone on HaskellCafe.

-- |2-D infinite grid with O(1) lookup, modification, and incremental move
-- 
-- Data is saved sparsely (wrapped in Maybe) with a rose tree zipper
-- where depth is manhattan distance from origin, and sibling index is order
-- CCW around a diamond centered at origin. Sparsity maximized by independent
-- angular offset for each radius, with only adjacent points stored.
--
-- Uses "Data.Tree.Zipper" which can be found on hackage at
--   
-- 
-- Data.Tree is indexed by Int. For unbounded grid, rewrite using Integer
--
-- Copyright (c) Dan Weston, 2008. All rights reserved.
--
-- License: Simplified BSD.

module GridZipper (
 -- * Grid representation
 module Data.Tree,
 module Data.Tree.Zipper,
 GridLabel(..),
 Grid,
 GridZipper,
 newGrid,
 -- * Grid coordinates
 XY(..),
 RK(..),
 getRK,getXY,
 cartesianFromPolar,polarFromCartesian,
 -- * Grid values
 getValue,newValue,setValue,
 -- * Moving around the grid
 goToRK,goToXY,moveInXY,north,south,east,west,
 -- * Example usage
 assocList,sampleUsage
 ) where


import Data.Tree.Zipper(TreeLoc,getLabel,setLabel,modifyLabel,
                        root,parent,left,right,firstChild,toTree,fromTree,
                        insertRight,insertDownFirst)
import Data.Tree       (Tree,flatten)
import Data.Maybe      (maybe,isJust,fromJust)

------------------------------------------------------------------------------
-- DATA TYPES
------------------------------------------------------------------------------

-- |Cartesian grid coordinates
data XY = XY Int Int deriving (Eq,Show)

-- |Polar grid coordinates
-- r = |x| + |y| (manhattan distance form origin)
-- k = index around diamond, starting at (+r,0)
data RK = RK Int Int deriving  (Eq,Show)

-- |Grid label
data GridLabel  a = GridLabel RK (Maybe a) deriving (Eq,Show)

-- |Grid represented as rose tree (radius = depth, angle = width)
type Grid       a = Tree    (GridLabel a)

-- |Cursor is rose tree zipper (polar coords stored in label alongside value)
type GridZipper a = TreeLoc (GridLabel a)


------------------------------------------------------------------------------
-- COORDINATE CONVERSION
------------------------------------------------------------------------------

-- |Gets cartesian coordinates from polar ones
cartesianFromPolar :: RK -> XY
cartesianFromPolar (RK 0 0) = XY 0 0
cartesianFromPolar (RK r k) = case q of
      0 -> XY (r - s   ) (s       )
      1 -> XY (negate s) (r - s   )
      2 -> XY (s - r   ) (negate s)
      3 -> XY (s       ) (s - r   )
  where (q,s) = k `divMod` r

-- |Gets polar coordinates from cartesian ones
polarFromCartesian :: XY -> RK
polarFromCartesian (XY 0 0) = RK 0 0
polarFromCartesian (XY x y)
  | x > 0 && y >= 0 = RK r        y
  | y > 0 && x <= 0 = RK r (r   - x)
  | x < 0 && y <= 0 = RK r (2*r - y)
  | y < 0 && x >= 0 = RK r (3*r + x)
  where r  = abs x + abs y

------------------------------------------------------------------------------
-- COORDINATE ACCESS
------------------------------------------------------------------------------

-- |Extracts polar coordinates from label
getRK :: GridLabel a -> RK
getRK (GridLabel rk _) = rk

-- |Extracts cartesian coordinates from label
getXY :: GridLabel a -> XY
getXY = cartesianFromPolar . getRK

------------------------------------------------------------------------------
-- VALUE ACCESS AND MODIFY
------------------------------------------------------------------------------

-- |Extracts grid value, if any, from label
getValue :: GridLabel a -> Maybe a
getValue (GridLabel _ value) = value

-- |Returns copy, replacing grid value
newValue :: Maybe a -> GridLabel a -> GridLabel a
newValue v (GridLabel rk _) = GridLabel rk v

-- |Returns copy, replacing grid value
setValue :: Maybe a -> GridZipper a -> GridZipper a
setValue v = modifyLabel (newValue v)

------------------------------------------------------------------------------
-- NODE CREATION
------------------------------------------------------------------------------

-- |New empty grid
newGrid :: Grid a
newGrid = newNode (RK 0 0)

------------------------------------------------------------------------------
-- MOVING THROUGH GRID
------------------------------------------------------------------------------

-- |Move to new polar coords
goToRK :: RK -> GridZipper a -> GridZipper a
goToRK rk@(RK r k) z
  | r <  0         = error "goToRK called with r < 0"
  | r == 0         = root z
  | r == rCurr     = moveAround    rk . leftmostSibling $ z
  | r >  rCurr     = moveOut rCurr rk z
  | otherwise      = moveIn  rCurr rk z
  where RK rCurr _ = getRK . getLabel $ z

-- Move to new cartesian coordinate
goToXY :: XY -> GridZipper a -> GridZipper a
goToXY = goToRK . polarFromCartesian

-- |Move relatively in delta cartesian coordinates
moveInXY :: Int -> Int -> GridZipper a -> GridZipper a
moveInXY dx dy z      = goToXY (XY (xOld + dx) (yOld + dy)) $ z
  where  XY xOld yOld = getXY . getLabel $ z

-- |Move up one step
north :: GridZipper a -> GridZipper a
north = moveInXY 0 1

-- |Move down one step
south :: GridZipper a -> GridZipper a
south = moveInXY 0 (-1)

-- |Move right one step
east  :: GridZipper a -> GridZipper a
east  = moveInXY 1 0

-- |Move left one step
west  :: GridZipper a -> GridZipper a
west  = moveInXY (-1) 0

-- |Display grid as association list: (xy,getValue)
assocList :: GridZipper a -> [(XY,a)]
assocList = map (\l -> (getXY               $ l,
                        fromJust . getValue $ l))
          . filter (isJust . getValue)
          . flatten
          . toTree
          . root

-- |Example of walking grid from origin, setting values
-- 
-- > sampleUsage = putStrLn . show . (assocList &&& id) . walkGrid . fromTree
-- >             $ (newGrid :: Grid String)
-- >   where f &&& g  = \x -> (f x, g x)
-- >         f >>> g  = g . f
-- >         walkGrid =   east           >>> setValue (Just "XY 1 0")
-- >                  >>> north >>> west >>> setValue (Just "XY 0 1")
-- >                  >>> south          >>> setValue (Just "XY 0 0")
-- >                  >>> south          >>> setValue (Just "XY 0 (-1)")
-- 
sampleUsage :: IO ()
sampleUsage = putStrLn . show . (assocList &&& id) . walkGrid . fromTree
     $ (newGrid :: Grid String)
  where f &&& g  = \x -> (f x, g x)
        f >>> g  = g . f
        walkGrid =   east           >>> setValue (Just "XY 1 0")
                 >>> north >>> west >>> setValue (Just "XY 0 1")
                 >>> south          >>> setValue (Just "XY 0 0")
                 >>> south          >>> setValue (Just "d(XY 0 (-1)")

------------------------------------------------------------------------------
-- HELPER FUNCTIONS NOT EXPORTED
------------------------------------------------------------------------------

-- |Returns a new node, intended for a given polar coordinate
-- Note that all grids are anchored at the origin. Only the origin node
-- functions as a valid standalone grid.
newNode :: RK -> Grid a
newNode rk  = return (GridLabel rk Nothing)

-- |Gets leftmost sibling of current node (which may be current one)
leftmostSibling :: GridZipper a -> GridZipper a
leftmostSibling z = maybe z leftmostSibling . left $ z

-- |Gets rightmost sibling of current node (which may be current one)
rightmostSibling :: GridZipper a -> GridZipper a
rightmostSibling z = maybe z rightmostSibling . right $ z

-- |Move inward to new polar coordinate
moveIn :: Int -> RK -> GridZipper a -> GridZipper a
moveIn rCurr rk@(RK r k) z
  | rCurr == r = moveAround         rk . leftmostSibling   $ z
  | otherwise  = moveIn (rCurr - 1) rk . fromJust . parent $ z

-- |Move outward to new polar coordinate
moveOut :: Int -> RK -> GridZipper a -> GridZipper a
moveOut rCurr rk@(RK r k) z
  | r == rCurr+1 = zChild
  | otherwise    = moveOut (rCurr + 1) rk zChild
  where zChild   = moveOutOne rk z

-- |Move outward exactly one unit of radius to new polar coordinate
-- This special case allows us to check if there is no child there and,
-- if so, to pick the angular anchor
-- Note that r passed in must be exactly one more than that of current node
moveOutOne :: RK -> GridZipper a -> GridZipper a
moveOutOne rk@(RK r k) z
  = maybe (insertDownFirst (newNode rk) z) (moveAround rk) $ firstChild z

-- |Move relatively in angle around origin (along diamond perimeter)
-- Note that r passed in must match that of current node
moveAround :: RK -> GridZipper a -> GridZipper a
moveAround rk@(RK r k) z
  | k == kCurr = z
  | otherwise  = maybe (insertRight (newNode rk) z) (moveAround rk) $ right z
  where RK _ kCurr = getRK . getLabel $ z