Axes.hs

{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE StandaloneDeriving #-}

data Axis = XAxis | YAxis

type family OtherAxis a where
  OtherAxis XAxis = YAxis
  OtherAxis YAxis = XAxis

class AxisCoords (a :: Axis) where
  data Coords a :: * -> *
  fromCoords :: (c, c) -> Coords a c

instance AxisCoords XAxis where
  data Coords XAxis a = XCoords a a
  fromCoords = uncurry XCoords

instance AxisCoords YAxis where
  data Coords YAxis a = YCoords a a
  fromCoords = uncurry YCoords

deriving instance Show a => Show (Coords XAxis a)
deriving instance Show a => Show (Coords YAxis a)
deriving instance Eq a => Eq (Coords XAxis a)
deriving instance Eq a => Eq (Coords YAxis a)
deriving instance Ord a => Ord (Coords XAxis a)

instance Ord a => Ord (Coords YAxis a) where
  compare (YCoords x1 y1) (YCoords x2 y2) = compare y1 y2 <> compare x1 x2

data Tree2D (x :: Axis) a b = Tree2D
  { treeCoords :: Coords x a
  , treeVal :: b
  , treeLeft :: Maybe (Tree2D (OtherAxis x) a b)
  , treeRight :: Maybe (Tree2D (OtherAxis x) a b) }

deriving instance (Show a, Show b) => Show (Tree2D XAxis a b)
deriving instance (Show a, Show b) => Show (Tree2D YAxis a b)

recurseTree2D :: (Ord (Coords x a), Ord (Coords (OtherAxis x) a), AxisCoords x, x ~ OtherAxis (OtherAxis x))
              => Tree2D x a b
              -> (a, a)
              -> Either (Tree2D x a b) (Maybe (Tree2D (OtherAxis x) a b))
recurseTree2D t@Tree2D { treeCoords = c, treeLeft = l, treeRight = r } key =
  case compare (fromCoords key) c of
    EQ -> Left t
    LT -> Right l
    GT -> Right r