Disjoint set in Haskell

Question

I think my error-handling here is atrocious, but I'm really new to Haskell so hoping one of you sages can help me out. This program is a command-line tool which you use by typing one of two commands: "Add " or "Check ". When you Add a pair of ints, it makes them connected, and when you Check two ints, it tells you whether they are connected (directly or indirectly). This uses the Union find data structure.

import Data.Map (Map)
import qualified Data.Map as Map
import qualified Text.Read as Read

data DisjointSet k = DisjointSet
  { parent :: Map k k
  , rank :: Map k Int
  }

empty :: (Ord k) => DisjointSet k
empty = DisjointSet {parent = Map.empty, rank = Map.empty}

make :: Ord k => k -> DisjointSet k -> DisjointSet k
make k d
  | (Map.member k (parent d)) = d
  | otherwise =
    d {parent = Map.insert k k (parent d), rank = Map.insert k 0 (rank d)}

find :: Ord k => k -> DisjointSet k -> Maybe (k, DisjointSet k)
find k d = do
  k' <- (Map.lookup k (parent d))
  if (k' == k)
    then return (k, d)
    else do
      (k'', d') <- find k' d
      return (k'', d' {parent = Map.insert k k'' (parent d')})

merge :: Ord k => k -> k -> DisjointSet k -> Maybe (DisjointSet k)
merge k1 k2 d =
  let d' = (make k1 (make k2 d))
  in do (k1', d'') <- find k1 d'
        (k2', d''') <- find k2 d''
        r1 <- Map.lookup k1' (rank d''')
        r2 <- Map.lookup k2' (rank d''')
        case compare r1 r2 of
          LT -> return d''' {parent = Map.insert k1' k2' (parent d''')}
          GT -> return d''' {parent = Map.insert k2' k1' (parent d''')}
          EQ ->
            return
              d'''
              { parent = Map.insert k1' k2' (parent d''')
              , rank = Map.insert k2' (r2 + 1) (rank d''')
              }

data Action
  = Add Int
        Int
  | Check Int
          Int
  deriving (Show, Read)

main :: IO ()
main = repl empty

repl :: DisjointSet Int -> IO ()
repl d = do
  line <- getLine
  case (Read.readMaybe line) of
    Just action ->
      case action of
        Add k1 k2 ->
          case (merge k1 k2 d) of
            Just d' -> repl d'
        Check k1 k2 ->
          case (find k1 d) of
            Just (k1', d') ->
              case (find k2 d') of
                Just (k2', d'') -> do
                  if (k1' == k2')
                    then putStrLn "Together"
                    else putStrLn "Apart"
                  repl d''

How do I clean up my code?

Answer 1

Naming

While it's usually not a big deal to have names like d and d' flying around, seeing d''' made me wonder: is there no better name for that?

Indentation

data Action
  = Add Int
        Int
  | Check Int
          Int
  deriving (Show, Read)

That seems.. unusal. Why is the second Int on a new line? Nice of you to seperate the constructors to their own lines, though.

data Action
    = Add Int Int
    | Check Int Int
    deriving (Show, Read)

Handling errors

You've said your error handling is atrocious. I wouldn't really call it that, because it simply does not exist. You never take any action on Nothing, which means the repl will stop as soon as something bad happens.

The glaring problem: The rat's tail

repl :: DisjointSet Int -> IO ()
repl d = do
  line <- getLine
  case (Read.readMaybe line) of
    Just action ->
      case action of
        Add k1 k2 ->
          case (merge k1 k2 d) of
            Just d' -> repl d'
        Check k1 k2 ->
          case (find k1 d) of
            Just (k1', d') ->
              case (find k2 d') of
                Just (k2', d'') -> do
                  if (k1' == k2')
                    then putStrLn "Together"
                    else putStrLn "Apart"
                  repl d''

Having lots of nested cases always yells: refactor me.

Let's follow the flow of repl for a second...

Get a line and read it to an Action. If that is possible, diverge on the type of action: if it's an Add, merge them and repeat on success. If it's a Check, find the first one. If it's found, find the second one. If that is found, show whether they are together or apart and repeat.

There are lots of actions and processes stuffed together here, so let's help the poor repl breathe by splitting it up into pieces.

Reading input

What should happen when the user types garbage, such as "Bananas are a fruit"? Maybe we should just try again... so let's do that.

getAction :: IO Action
getAction = case Read.readMaybe <$> getLine of
    Just action -> return action
    Nothing     -> putStrLn "Oops, try again..." >> getAction

We read a line of user input and Read.readMaybe it to form an Action (<$> is just the infix operator for fmap). If that succeeds, horray, we got ourselves an Action! If not.. say something to the user and try again.

Executing the action

What should happen on an error here? That is your decision to make. If anything goes wrong, I'll just return the old DisjointSet and do nothing. I like to indent the case-statements based on the indentation level on the case itself, so let's do that.

execute :: DisjointSet Int -> Action -> IO (DisjointSet Int)
execute d (Add a b)   = return $ maybe d id (merge a b d)
execute d (Check a b) = case find a d of
                          Nothing       -> return d
                          Just (a', d') -> case find b d' of
                                             Nothing         -> return d
                                             Just (b', d'') -> if a' == b'
                                                                then putStrLn "Together" >> return d''
                                                                else putStrln "Apart"    >> return d''

(Note to Editors: The long line is on purpose) Damn, that's still long.

Monads to the rescue!

The beauty of Maybe as a monad is that it'll stop doing whatever it does and just return Nothing once it encounters one. So let's use that!

execute :: DisjointSet Int -> Action -> IO (Maybe (DisjointSet Int))
execute d (Add a b)   = return $ merge a b d
execute d (Check a b) = do
    (together, d'') <- (do
             (a', d')  <- find a d
             (b', d'') <- find b d'
             return (a' == b', d''))
    when together       (putStrLn "Together")
    when (not together) (putStrLn "Apart")
    return d''

Putting it together

Using these fancy new functions, your new repl would be:

repl :: DisjointSet Int -> IO ()
repl d = getAction >>= execute d >>= \c -> case c of
    Just d' -> repl d'
    Nothing -> return ()

Answer 2

Control.Lens makes the record syntax pain go away. non unmakes make, and eliminates the worry about missing parents. That said, we never actually profit from initializing nodes with themselves as parents. State abstracts away the passery. untilJust pushes the worry about malformed input off into infinity. Control.Monad.Loops should provide iterateMaybeM, but I think my name is stupid and my implementation might be stricter than needed in the general case, so I don't want to submit that PR. I'll cease caching find in parent to make find a one-liner.

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TemplateHaskell #-}
{-# Language LambdaCase #-}
import Data.Map (Map)
import qualified Data.Map as Map
import qualified Text.Read as Read
import Control.Monad.IO.Class (liftIO)
import Control.Monad.State
import Data.Foldable (asum)
import Control.Lens
import Control.Monad.Loops (untilJust, unfoldrM)

data DisjointSet k = DisjointSet
  { _parent :: Map k k
  , _sparseRank :: Map k Int
  }
makeLenses ''DisjointSet

rank k = sparseRank . at k . non 0

iterateMaybeM :: Monad m => (a -> m (Maybe a)) -> a -> m [a]
iterateMaybeM f = unfoldrM $ (fmap . fmap) (\a -> (a, a)) . f

find :: (MonadState (DisjointSet k) m, Ord k) => k -> m k
find k = last <$> iterateMaybeM (\k -> use $ parent . at k) k

data Action = Add Int Int | Check Int Int deriving (Show, Read)

main :: IO ()
main = (`evalStateT` DisjointSet Map.empty Map.empty) $ forever $
  untilJust (Read.readMaybe <$> liftIO getLine) >>= \case
    Add k1 k2 -> do
      r1 <- find k1
      r2 <- find k2
      compare <$> use (rank r1) <*> use (rank r2) >>= \case
        LT -> parent . at r1 ?= r2
        GT -> parent . at r2 ?= r1
        EQ -> do
          parent . at r1 ?= r2
          rank r2 += 1
    Check k1 k2 -> do
      together <- (==) <$> find k1 <*> find k2
      liftIO $ putStrLn $ if together then "Together" else "Apart"