
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