Skip to content

Instantly share code, notes, and snippets.

@FranklinChen

FranklinChen/walk-tree.hs

Created April 26, 2015 03:09
Show Gist options
  • Save FranklinChen/1f43aa23ff87b514a265 to your computer and use it in GitHub Desktop.
Save FranklinChen/1f43aa23ff87b514a265 to your computer and use it in GitHub Desktop.
Download ZIP
Merge a list of annotations into selected nodes of a tree, with error recovery and reporting
Raw
walk-tree.hs
-- | Something for work, prototyped first in Haskell before turning
-- into Scala
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE LambdaCase #-}
import Test.Hspec
import Control.Monad.State
-- | A tree of values.
data Content a = Leaf a
| Node [Content a]
deriving (Show, Eq, Functor, Foldable, Traversable)
-- | Annotating a value with a label.
data Annotated label a = Annotated label a
deriving (Show, Eq, Functor, Foldable, Traversable)
-- | Pair every 'Right' with corresponding annotation.
-- If too many labels, we still finish, always returning the leftover labels.
-- If not enough labels, we annotate with 'Nothing', so we must use
-- 'Maybe'.
merge :: Content (Either a b)
-> [label]
-> (Content (Either a (Annotated (Maybe label) b)), [label])
merge = runState . traverse mergeLeafS
-- | Attempt to merge a leaf, using supply of labels
-- Error recovery for leftover content with no labels is to use
-- 'Nothing'. This means when there are labels we must use 'Just'.
mergeLeafS :: Either a b
-> State [label] (Either a (Annotated (Maybe label) b))
mergeLeafS (Left l) = pure (Left l)
mergeLeafS (Right r) = get >>= \case
[] -> pure (Right $ Annotated Nothing r) -- error recovery
label : labels ->
put labels >>
pure (Right $ Annotated (Just label) r)
-- | Grab just the 'Nothing' labeled nodes.
contentsMissingLabels :: Content (Either a (Annotated (Maybe label) b))
-> [b]
contentsMissingLabels = foldr combine [] where
combine (Right (Annotated Nothing b)) bs = b : bs
combine _ bs = bs
-- | HSpec tests.
main :: IO ()
main = hspec spec
spec :: Spec
spec = describe "walk tree" $ do
describe "merge" $ do
it "handles matched lengths" $ do
merge contents1 [1..3] `shouldBe` (enoughContents1, [])
it "handles leftover labels" $ do
merge contents1 [1..5] `shouldBe` (enoughContents1, [4, 5])
it "handles not enough labels" $ do
merge contents1 [1] `shouldBe` (notEnoughContents1, [])
describe "contents with not enough labels" $ do
it "finds no missing labels" $ do
contentsMissingLabels enoughContents1 `shouldBe` []
it "finds missing labels" $ do
contentsMissingLabels notEnoughContents1 `shouldBe` ["def", "ghi"]
contents1 :: Content (Either Char String)
contents1 = Node [ Leaf $ Left 'a'
, Leaf $ Right "abc"
, Node [ Leaf $ Left 'b'
, Leaf $ Right "def"
, Leaf $ Left 'c'
]
, Leaf $ Right "ghi"
]
enoughContents1 :: Content (Either Char (Annotated (Maybe Int) String))
enoughContents1 = Node [ Leaf $ Left 'a'
, Leaf $ Right $ Annotated (Just 1) "abc"
, Node [ Leaf $ Left 'b'
, Leaf $ Right $ Annotated (Just 2) "def"
, Leaf $ Left 'c'
]
, Leaf $ Right $ Annotated (Just 3) "ghi"
]
notEnoughContents1 :: Content (Either Char (Annotated (Maybe Int) String))
notEnoughContents1 = Node [ Leaf $ Left 'a'
, Leaf $ Right $ Annotated (Just 1) "abc"
, Node [ Leaf $ Left 'b'
, Leaf $ Right $ Annotated Nothing "def"
, Leaf $ Left 'c'
]
, Leaf $ Right $ Annotated Nothing "ghi"
]
@FranklinChen
Copy link
Author

Of course, it is not actually necessary to use State here. But the more general case involving logging may want a monad anyway. For the exact example here, tuples aren't too bad:

import Data.Traversable (mapAccumL)

-- | Without using 'State'.
mergeDirect :: [label]
            -> Content (Either a b)
            -> ([label], Content (Either a (Annotated (Maybe label) b)))
mergeDirect = mapAccumL mergeLeaf

-- | For use with mapAccumL
mergeLeaf :: [label]
          -> Either a b
          -> ([label], Either a (Annotated (Maybe label) b))
mergeLeaf labels (Left l) = (labels, Left l)
mergeLeaf [] (Right r) = ([], Right $ Annotated Nothing r)
mergeLeaf (label : labels) (Right r) = (labels, Right $ Annotated (Just label) r)

@FranklinChen
Copy link
Author

Scala version:

import scalaz._
import Scalaz._

object WalkTree {
  /** Tree. */
  sealed trait Content[+A]

  final case class Leaf[+A](leaf: A) extends Content[A]
  final case class Node[+A](nodes: List[Content[A]]) extends Content[A]

  final case class Annotated[+Label, +A](label: Label, a: A)

  implicit def treeInstances: Traverse[Content] =
    new Traverse[Content] {
      import scala.language.higherKinds

      override def traverseImpl[G[_]: Applicative, A, B]
        (fa: Content[A])(f: A => G[B]):
          G[Content[B]] = fa match {
        case Leaf(leaf) => f(leaf).map(Leaf.apply)
        case Node(nodes) => nodes.traverse(_.traverse(f)).map(Node.apply)
    }
  }

  /**
    Use special Scalaz support for traversing with state and trampolining.
    */
  def mergeLabels[A, B, Label](
    content: Content[A \/ B],
    labels: List[Label]
  ): (List[Label], Content[A \/ Annotated[Option[Label], B]]) =
    content.runTraverseS(labels)(mergeLeafS)

  def mergeLeafS[A, B, Label](leaf: A \/ B):
      State[List[Label], A \/ Annotated[Option[Label], B]] = {
    val S = StateT.stateMonad[List[Label]]
    import S.monadSyntax._

    leaf match {
      case -\/(l) => pure(-\/(l))
      case \/-(r) => get[List[Label]] >>= {
        case Nil => pure(\/-(Annotated(None, r)))
        case label::labels => put(labels) >>
          pure(\/-(Annotated(some(label), r)))
      }
    }
  }
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment