[Haskell-cafe] Re: Red-Blue Stack

[Timothy Goddard]

There was a bug in there with popping the non-head colour off the stack. 
Updated code, please test thoroughly:

module RBStack where

data RBColour = Red | Blue
  deriving (Show, Eq)

data RBStack a =  RBStack {
  headColour :: RBColour,
  stackElems :: [[a]]
}
  deriving (Show, Eq)

otherCol :: RBColour -> RBColour
otherCol Red = Blue
otherCol Blue = Red

empty :: RBStack a
empty = RBStack Red []

push :: RBColour -> a -> RBStack a -> RBStack a
push col val stack
  | null (stackElems stack) = RBStack col [[val]]
  | headColour stack == col = RBStack col ((val:e):es)
  | otherwise = RBStack col ([val]:e:es)
  where
    (e:es) = stackElems stack

popColour :: RBColour -> RBStack a -> (Maybe a, RBStack a)
popColour col stack
  | null (stackElems stack) = (Nothing, stack)
  | headColour stack == col = (Just (head e), if null (tail e)
      then (RBStack (otherCol col) es)
      else (RBStack col ((tail e):es)))
  | otherwise = if null es
      then (Nothing, stack)
      else let (f:fs) = es in (Just (head f), if null (tail f)
        then (if null fs then (RBStack (otherCol col) [e]) else (RBStack 
(otherCol col) ((e ++ (head fs)):(tail fs))))
        else RBStack (otherCol col) (e:(tail f):fs))
  where
    (e:es) = stackElems stack

pop :: RBStack a -> (Maybe (RBColour, a), RBStack a)
pop stack
  | null (stackElems stack) = (Nothing, stack)
  | otherwise = (Just (col, head e), if null (tail e) then (RBStack (otherCol 
col) es) else (RBStack col ((tail e):es)))
  where
    (e:es) = stackElems stack
    col = headColour stack

peek :: RBStack a -> Maybe (RBColour, a)
peek = fst . pop