How to fix this Haskell lookup function type mismatch?

Question

If someone could lend a hand fixing this error, it would be much appreciated. The code is:

type Name = String
type Coordinates = (Int, Int)
type Pop = Int
type TotalPop = [Pop]
type City = (Name, (Coordinates, TotalPop))

testData :: [City]
testData = [("New York City", ((1,1), [5, 4, 3, 2])),
           ("Washingotn DC", ((3,3), [3, 2, 1, 1])),
           ("Los Angeles", ((2,2), [7, 7, 7, 5]))]

getCityPopulation :: [City] -> Name -> Int -> Int
getCityPopulation cs nameIn yearIn = head ([ z !! (yearIn - 1) | (x,z) <- lookup nameIn cs])

Expected behaviour:

getCityPopulation testData "New York City" 2
>>> 4

Actual behaviour:

    • Couldn't match expected type ‘[(a0, [Int])]’
                  with actual type ‘Maybe (Coordinates, TotalPop)’
    • In the expression: lookup nameIn cs
      In a stmt of a list comprehension: (x, z) <- lookup nameIn cs
      In the first argument of ‘head’, namely
        ‘([z !! (yearIn - 1) | (x, z) <- lookup nameIn cs])’
   |
55 | getCityPopulation cs nameIn yearIn = head ([ z !! (yearIn - 1) | (x,z) <- lookup nameIn cs])
   |                                                                           ^^^^^^^^^^^^^^^^

Answer 1

getCityPopulation :: [City] -> Name -> Int -> Int
getCityPopulation [] nameIn yearIn = error "Can't be empty"
getCityPopulation cs nameIn yearIn = 
   head ([ z !! (yearIn - 1) | (x,z) <- maybeToList $ lookup nameIn cs])

does the job. This uses

maybeToList :: Maybe a -> [a]   -- Defined in `Data.Maybe'

which is needed to transform the output of

lookup :: Eq a => a -> [(a, b)] -> Maybe b

Of course that code should actually be

   head ([ z !! (yearIn - 1) | (x,z) <- maybeToList $ lookup nameIn cs])
 ==
   case (lookup nameIn cs) of 
      Just (x,z) ->  z !! (yearIn - 1)
      Nothing    ->  error "couldn't find it"

so that

> getCityPopulation testData "New York City" 2
4

> getCityPopulation testData "Las Vegas" 2
*** Exception: couldn't find it

but writing a code that can fail like that is also not right. It is better to again wrap the result in a Maybe:

getCityPopulation :: [City] -> Name -> Int -> Maybe Int
getCityPopulation [] nameIn yearIn = error "Can't be empty"
getCityPopulation cs nameIn yearIn = 
   case (lookup nameIn cs) of 
      Just (x,z) ->  Just $ z !! (yearIn - 1)
      Nothing    ->  Nothing

and then we have

> getCityPopulation testData "New York City" 2
Just 4

> getCityPopulation testData "Las Vegas" 2
Nothing

Now, having changed the output type, we can actually go back to the original code -- almost:

getCityPopulation :: [City] -> Name -> Int -> Maybe Int
getCityPopulation [] nameIn yearIn = error "Can't be empty"
getCityPopulation cs nameIn yearIn = 
   [ z !! (yearIn - 1) | (x,z) <- lookup nameIn cs]

What magic is this, you ask? It is known as MonadComprehensions. You turn it on either at GHCi prompt with

> :set -XMonadComprehensions

or by including

{-# LANGUAGE MonadComprehensions #-}

pragma at the top of your source file.

There's no magic fix though for the possible out-of-bounds access problem that your code has, using that !! without checking. Or is there?

getCityPopulation :: [City] -> Name -> Int -> Maybe Int
getCityPopulation [] nameIn yearIn = error "Can't be empty"
getCityPopulation cs nameIn yearIn = 
   [ e | (_,z) <- lookup nameIn cs, 
         e <- listToMaybe $ drop (yearIn-1) z ]