AdventOfCode 2019 day 6 in Haskell

Question

I am new to Haskell and currently trying to port my solutions for the 2019 installment of the coding challenge AdventOfCode to Haskell. So, I would very much appreciate any suggestions how to make the code more readable and, in particular, more idiomatic.

This post shows my solution of day 6 part 2, but also includes the function totalDecendantCount used to solve part 1. If you have not solved these problems and still intend to do so, stop reading immediately.

For both problems, you get a file with an orbit specification on each line of the form A)B, which tells you that B orbits A. This describes a tree of bodies orbiting each other with root COM.

In part 1, you have to compute a check sum. More precisely, you have to compute the sum of the number of direct and indirect orbits of each body, which is the same as the sum of the number of descendants of each body in the tree.

In part 2, which you cannot see if you have not finished part 1, you have to compute the minimal number of transfers between orbits from you (YOU) to Santa (SAN).

I have kept the entire solution for each part of each day in a single module with a single exported function that prints the solution. For day 6 part 2 it starts as follows.

module AdventOfCode20191206_2
 ( distanceToSanta
 ) where
import System.IO
import Data.List.Split
import Data.List
import Data.Maybe
import Data.Hashable
import qualified Data.HashMap.Strict as Map
distanceToSanta :: IO ()
distanceToSanta = do
 inputText <- readFile "Advent20191206_1_input.txt"
 let orbitList = (map orbit . lines) inputText
 let orbits = orbitMap $ catMaybes orbitList
 let pathToSanta = fromJust $ path orbits "COM" "YOU" "SAN"
 let requiredTransfers = length pathToSanta - 3
 print requiredTransfers

We subtract 3 from the length of the path because it consists of the bodies on the path and you only have to transfer from the body you already orbit to the body Santa orbits.

To store the tree, I use a HashMap.Strict and introduce the following type aliases and helper function to make things a bit more descriptive.

type OrbitSpecification = (String,String)
type ChildrenMap a = Map.HashMap a [a]
children :: (Eq a, Hashable a) => ChildrenMap a -> a -> [a]
children childrenMap = fromMaybe [] . flip Map.lookup childrenMap

Next follow the functions I use to read in the tree.

orbit :: String -> Maybe OrbitSpecification
orbit str =
 case orbit_specification of
 [x,y] -> Just (x,y)
 _ -> Nothing
 where orbit_specification = splitOn ")" str
orbitMap :: [OrbitSpecification] -> ChildrenMap String
orbitMap = Map.fromListWith (++) . map (applyToSecondElement toSingleElementList)
applyToSecondElement :: (b -> c) -> (a,b) -> (a,c)
applyToSecondElement f (x,y) = (x, f y)
toSingleElementList :: a -> [a]
toSingleElementList x = [x]

To solve part 1, I introduce two general helper function to generate aggregates over children or over all descendents.

childrenAggregate :: (Eq a, Hashable a) => ([a] -> b) -> ChildrenMap a -> a -> b
childrenAggregate aggregatorFnc childrenMap = aggregatorFnc . children childrenMap
decendantAggregate :: (Eq a, Hashable a) => (b -> b -> b) -> (ChildrenMap a -> a -> b) -> ChildrenMap a -> a -> b
decendantAggregate resultFoldFnc nodeFnc childrenMap node =
 foldl' resultFoldFnc nodeValue childResults
 where
 nodeValue = nodeFnc childrenMap node
 childFnc = decendantAggregate resultFoldFnc nodeFnc childrenMap
 childResults = map childFnc $ children childrenMap node

The descendantAggragate recursively applies a function nodeFnc to a node node and all its descendants and folds the results using some function resultFoldFnc. This allows to define the necessary functions to count the total number of descendants of a node as follows.

childrenCount :: (Eq a, Hashable a) => ChildrenMap a -> a -> Int
childrenCount = childrenAggregate length
decendantCount :: (Eq a, Hashable a) => ChildrenMap a -> a -> Int
decendantCount = decendantAggregate (+) childrenCount
totalDecendantCount :: (Eq a, Hashable a) => ChildrenMap a -> a -> Int
totalDecendantCount = decendantAggregate (+) decendantCount

For part 2, we use that between two points in a tree, there is exactly one path (without repetition). First, we define a function to get a path from the root of a (sub)tree to the destination, provided it exists.

pathFromRoot :: (Eq a, Hashable a) => ChildrenMap a -> a -> a -> Maybe [a]
pathFromRoot childrenMap root destination
 | destination == root = Just [root]
 | null childPaths = Nothing
 | otherwise = Just $ root:(head childPaths)
 where
 rootChildren = children childrenMap root
 pathFromNewRoot newRoot = pathFromRoot childrenMap newRoot destination
 childPaths = mapMaybe pathFromNewRoot rootChildren

This function only finds paths down from the root of a (sub)tree. General paths come in three variations: path from the root of a (sub)tree, the inverse of such a path or the concatenation of a path to the root of a subtree and one from that root to the end point. Thus, we get the path as follows.

path :: (Eq a, Hashable a) => ChildrenMap a -> a -> a -> a -> Maybe [a]
path childrenMap root start end =
 let maybeStartEndPath = pathFromRoot childrenMap start end
 in if isJust maybeStartEndPath
 then maybeStartEndPath
 else let maybeEndStartPath = pathFromRoot childrenMap end start
 in case maybeEndStartPath of
 Just endStartPath -> Just $ reverse endStartPath
 Nothing -> let
 rootPathToStart = pathFromRoot childrenMap root start
 rootPathToEnd = pathFromRoot childrenMap root end
 in if isNothing rootPathToStart || isNothing rootPathToEnd
 then Nothing
 else connectedPath (fromJust rootPathToStart) (fromJust rootPathToEnd)

To connect the paths in the last alternative, we follow both paths from the root to the last common point and then build it by concatenation the reverse of the path to the start with the path to the destination.

connectedPath :: Eq a => [a] -> [a] -> Maybe [a]
connectedPath rootToStart rootToEnd =
 case pathPieces of
 Nothing -> Nothing
 Just (middle, middleToStart, middleToEnd) ->
 Just $ (reverse middleToStart) ++ [middle] ++ middleToEnd
 where pathPieces = distinctPathPieces rootToStart rootToEnd
distinctPathPieces :: Eq a => [a] -> [a] -> Maybe (a, [a], [a])
distinctPathPieces [x] [y] = if x == y then Just (x, [], []) else Nothing
distinctPathPieces (x1:y1:z1) (x2:y2:z2)
 | x1 /= x2 = Nothing
 | y1 /= y2 = Just (x1, y1:z1, y2:z2)
 | otherwise = distinctPathPieces (y1:z1) (y2:z2)
distinctPathPieces _ _ = Nothing

This solution heavily depends on the input describing a tree. In case a DAG is provided, a result will be produced that is not necessary correct. For totalDescendantCount, nodes after joining branches will be counted multiple times and path will find a path, but not nescesarily the shortest one. If there are cycles in the graph provided, the recursions in the functions will not terminate.

Answer 1

Simplification

In path, notice how the code gets more nested as you try each possible path (either from start to end, or end to start, or from end to root and root to start). You can use the Alternative instance for Maybe to simplify this code:

let maybeStartEndPath = pathFromRoot childrenMap start end
 maybeEndStartPath = pathFromRoot childrenMap end start
 maybeRootPath = [...] -- see below
in 
 maybeStartEndPath
 <|> fmap reverse maybeEndStartPath
 <|> maybeRootPath

This code will try maybeStartEndPath first. If it returns Nothing, it will move on to the next option and so on.

For your final case (which I've named maybeRootPath), you do the following check:

if isNothing rootPathToStart || isNothing rootPathToEnd
 then Nothing
 else connectedPath (fromJust rootPathToStart) (fromJust rootPathToEnd)

This is more consicely done with liftA2 from Control.Applicative. liftA2 lifts a binary function into an applicative context:

λ :set -XTypeApplications
λ :t liftA2 @Maybe
liftA2 @Maybe :: (a -> b -> c) -> (Maybe a -> Maybe b -> Maybe c)

Then, if either argument is Nothing, the function will return Nothing without having to pattern match. So we can fill in maybeRootPath above with

 maybeRootPath = join $ liftA2 connectedPath rootPathToStart rootPathToEnd
 where
 rootPathToStart = pathFromRoot childrenMap root start
 rootPathToEnd = pathFromRoot childrenMap root end

The join is needed because connectedPath returns a Maybe already, and we've lifted it into Maybe, which leaves us with a return value of Maybe (Maybe [a]). join flattens nested monads, bringing us back to Maybe [a]

---

Minor points

Your function applyToSecondElement is second from Control.Arrow

λ :t second @(->)
second @(->) :: (b -> c) -> (d, b) -> (d, c)

---

toSingleElementList can also be written as (:[]) or return

So orbitMap can be written

orbitMap = Map.fromListWith (++) . map (second (:[]))

To your credit, your naming made both of these functions clear anyway, but it's more recognizable if you use functions that already exist.

---

Algorithm

I was going to suggest keeping each edge bidirectional instead of one-directional, so that you can directly check for a path from start to end instead of checking 3 cases. After reviewing the code, I think your approach is better from a functional perspective because it eliminates the need for you to check for cycles and keep a set as you search the graph. Good work.

---

Revised Code

import Control.Applicative
import Control.Monad
import Control.Arrow
import System.IO
import Data.List.Split
import Data.List
import Data.Maybe
import Data.Hashable
import qualified Data.HashMap.Strict as Map
main :: IO ()
main = do
 inputText <- readFile "Advent20191206_1_input.txt"
 let orbitList = catMaybes $ (map orbit . lines) inputText
 let orbits = orbitMap orbitList
 let pathToSanta = fromJust $ path orbits "COM" "YOU" "SAN"
 let requiredTransfers = length pathToSanta - 3
 print requiredTransfers
type OrbitSpecification = (String,String)
type ChildrenMap a = Map.HashMap a [a]
children :: (Eq a, Hashable a) => ChildrenMap a -> a -> [a]
children childrenMap = fromMaybe [] . flip Map.lookup childrenMap
orbit :: String -> Maybe OrbitSpecification
orbit str =
 case orbit_specification of
 [x,y] -> Just (x, y)
 _ -> Nothing
 where orbit_specification = splitOn ")" str
orbitMap :: [OrbitSpecification] -> ChildrenMap String
orbitMap = Map.fromListWith (++) . map (second (:[]))
childrenAggregate :: (Eq a, Hashable a) => ([a] -> b) -> ChildrenMap a -> a -> b
childrenAggregate aggregatorFnc childrenMap = aggregatorFnc . children childrenMap
decendantAggregate :: (Eq a, Hashable a) => (b -> b -> b) -> (ChildrenMap a -> a -> b) -> ChildrenMap a -> a -> b
decendantAggregate resultFoldFnc nodeFnc childrenMap node =
 foldl' resultFoldFnc nodeValue childResults
 where
 nodeValue = nodeFnc childrenMap node
 childFnc = decendantAggregate resultFoldFnc nodeFnc childrenMap
 childResults = map childFnc $ children childrenMap node
childrenCount :: (Eq a, Hashable a) => ChildrenMap a -> a -> Int
childrenCount = childrenAggregate length
decendantCount :: (Eq a, Hashable a) => ChildrenMap a -> a -> Int
decendantCount = decendantAggregate (+) childrenCount
totalDecendantCount :: (Eq a, Hashable a) => ChildrenMap a -> a -> Int
totalDecendantCount = decendantAggregate (+) decendantCount
pathFromRoot :: (Eq a, Hashable a) => ChildrenMap a -> a -> a -> Maybe [a]
pathFromRoot childrenMap root destination
 | destination == root = Just [root]
 | null childPaths = Nothing
 | otherwise = Just $ root:(head childPaths)
 where
 rootChildren = children childrenMap root
 pathFromNewRoot newRoot = pathFromRoot childrenMap newRoot destination
 childPaths = mapMaybe pathFromNewRoot rootChildren
path :: (Eq a, Hashable a) => ChildrenMap a -> a -> a -> a -> Maybe [a]
path childrenMap root start end =
 let maybeStartEndPath = pathFromRoot childrenMap start end
 maybeEndStartPath = pathFromRoot childrenMap end start
 maybeRootPath = join $ liftA2 connectedPath rootPathToStart rootPathToEnd
 where
 rootPathToStart = pathFromRoot childrenMap root start
 rootPathToEnd = pathFromRoot childrenMap root end
 in
 maybeStartEndPath
 <|> fmap reverse maybeEndStartPath
 <|> maybeRootPath
connectedPath :: Eq a => [a] -> [a] -> Maybe [a]
connectedPath rootToStart rootToEnd =
 case pathPieces of
 Nothing -> Nothing
 Just (middle, middleToStart, middleToEnd) ->
 Just $ (reverse middleToStart) ++ [middle] ++ middleToEnd
 where pathPieces = distinctPathPieces rootToStart rootToEnd
distinctPathPieces :: Eq a => [a] -> [a] -> Maybe (a, [a], [a])
distinctPathPieces [x] [y] = if x == y then Just (x, [], []) else Nothing
distinctPathPieces (x1:y1:z1) (x2:y2:z2)
 | x1 /= x2 = Nothing
 | y1 /= y2 = Just (x1, y1:z1, y2:z2)
 | otherwise = distinctPathPieces (y1:z1) (y2:z2)
distinctPathPieces _ _ = Nothing

• M.Doerner: Thank you very much! I already had the feeling that there had to be some function to apply a function to the second element, but did not know Arrow. I will need some time to look things up in order to fully understand your simplifications. But now I know some directions to look into to improve my Haskell.