A whirlwind tour of Haskell, day 7

Day 7, functionally: bit fiddling, Map, and recursion

Today's puzzle is quite interesting, and will allow me to demonstrate several Haskell techniques! We're given the description of an electronic circuit consisting of a bunch of constants and bitwise operations, and we need to simulate it. The bitwise operations are fairly standard, here's what they look like in Haskell:

import Data.Word
import Data.Bits

opAND :: Word16 -> Word16 -> Word16
opAND = (.&.)

opOR :: Word16 -> Word16 -> Word16
opOR = (.|.)

opLSHIFT :: Word16 -> Int -> Word16
opLSHIFT = shiftL

opRSHIFT :: Word16 -> Int -> Word16
opRSHIFT = shiftR

opNOT :: Word16 -> Word16
opNOT = complement

Like before, I'll skip the parsing details and I'll represent the operations using a precise type.

type WireName = String

data Expr
  = Number Word16
  | Wire   WireName
  | AND    Expr Expr
  | OR     Expr Expr
  | LSHIFT Expr Int
  | RSHIFT Expr Int
  | NOT    Expr
  deriving (Show, Eq)

-- (wire, expression producing its value)
type Assignment = (WireName, Expr)

Since numbers and wires are allowed both at the top-level (for example 123 -> a and w -> a) and nested inside a single operation (for example 1 OR b -> a), I found it slightly easier to allow operations to be nested into arbitrarily deep expressions. This way, I don't have to repeat my number and wire logic in two separate places.

Anyway, the tricky part of this puzzle is that those expressions aren't meant to be evaluated in the order in which they are given, as some of the wire names are used as input much earlier than the expressions which give a value to those wires. Instead, we should take our list of assignments and store them in a Map, so we'll be able to easily get the expression which produces the value of each wire.

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

type Env = Map WireName Expr

mkEnv :: [Assignment] -> Env
mkEnv = Map.fromList

lookupEnv :: Env -> WireName -> Maybe Expr
lookupEnv = flip Map.lookup

-- throws an exception if the name is not found
lookupEnv' :: Env -> WireName -> Expr
lookupEnv' = (Map.!)

Now that I can lookup the expression for each wire, I can evaluate expressions using a recursive solution: if the expression is a number (the base case), I'm done, if the expression depends on another wire, I lookup its expression and evaluate it using a recursive call, and if it's an operation, I recursively evaluate its operands and perform the operation. Easy!

eval :: Env -> Expr -> Word16
eval _   (Number   n) = n
eval env (Wire  name) = eval env (lookupEnv' env name)
eval env (AND    x y) = opAND    (eval env x) (eval env y)
eval env (OR     x y) = opOR     (eval env x) (eval env y)
eval env (LSHIFT x i) = opLSHIFT (eval env x) i
eval env (RSHIFT x i) = opRSHIFT (eval env x) i
eval env (NOT    x  ) = opNOT    (eval env x)

Using this solution, we can easily evaluate each of the wires given in the sample circuit:

-- |
-- >>> :{
-- let assignments = [ ("d", AND (Wire "x") (Wire "y"))
--                   , ("e", OR  (Wire "x") (Wire "y"))
--                   , ("f", LSHIFT (Wire "x") 2)
--                   , ("g", RSHIFT (Wire "y") 2)
--                   , ("h", NOT (Wire "x"))
--                   , ("i", NOT (Wire "y"))
--                   , ("x", Number 123)
--                   , ("y", Number 456)
--                   ]
--     wireNames   = map fst assignments
--     wireValues  = map (day7 assignments) wireNames
-- in forM_ (zip wireNames wireValues) print
-- :}
-- ("d",72)
-- ("e",507)
-- ("f",492)
-- ("g",114)
-- ("h",65412)
-- ("i",65079)
-- ("x",123)
-- ("y",456)
day7 :: [Assignment] -> WireName -> Word16
day7 assignments name = eval env (lookupEnv' env name)
  where
    env :: Env
    env = mkEnv assignments

However, if we try it on the puzzle's larger circuit, the algorithm gets stuck! What happens is that many of the wires are used more than once, and yet are re-evaluated each time they are encountered. Consider the expression b AND c -> a. If operands b and c both end up depending on the same wire d, evaluating a will end up evaluating d's expression twice, and if a has to be evaluated twice, d will end up being evaluated four times. Since the number of evaluations doubles each time a new element is added to such a chain, it doesn't take many elements for the number of evaluations to become astronomical, causing our algorithm to get stuck re-evaluating the same expressions over and over.

Day 7, imperatively: caching

Clearly, we should cache the result of each wire in order to avoid evaluating the same expression more than once. Here is an imperative implementation of that idea: keep a map of the results computed so far, and each time we encounter a wire name, look inside the map to see if we already have a result for this wire. If we do, use it, otherwise recursively compute the answer and store it in the map.

type Cache = Map WireName Word16

lookupCache :: Cache -> WireName -> Maybe Word16
lookupCache = flip Map.lookup

lookupCache' :: Cache -> WireName -> Word16
lookupCache' = (Map.!)

eval :: Env -> Expr -> State Cache Word16
eval _   (Number   n) = return n
eval env (Wire  name) = do
    cache <- get
    case lookupCache cache name of
      Just x  -> return x
      Nothing -> do
          x <- eval env (lookupEnv' env name)
          modify (Map.insert name x)
          return x
eval env (AND    x y) = opAND    <$> eval env x <*> eval env y
eval env (OR     x y) = opOR     <$> eval env x <*> eval env y
eval env (LSHIFT x i) = opLSHIFT <$> eval env x <*> pure i
eval env (RSHIFT x i) = opRSHIFT <$> eval env x <*> pure i
eval env (NOT    x  ) = opNOT    <$> eval env x

Day 7, lazily: laziness-based caching and MemoTrie

Haskell has another way to cache results which is built into the language: its implementation of laziness.

In the previous days, I pointed out a few situations in which laziness improves the performance by skipping over the part of a pure computation whose result does not end up being used. If the result does end up being used, the computation must be performed. If the result ends up being used multiple times, the computation must also be performed, but only once, not multiple times. The way in which this is implemented is that the memory in which the result would normally be stored can point to either a thunk describing which computation needs to be forced to obtain the result, or it can point to the result itself if it has already been computed. In other words, laziness caches the results of computations so that each computation only has to be performed at most once.

Another aspect of laziness is that it allows us to define infinite data structures, by only constructing the portion of the data which ends up being needed, and also circular data structures, because by the time the structure loops back to the beginning, the beginning's thunk has already been expanded into a value.

-- |
-- >>> take 10 circular
-- [1,2,3,1,2,3,1,2,3,1]
circular :: [Int]
circular = 1 : 2 : 3 : circular

Here is a trickier example in which instead of directly looping back to the beginning, we examine the portion of the data structure which has been constructed so far and we produce a modified version.

-- |
-- >>> take 10 nats
-- [1,2,3,4,5,6,7,8,9,10]
nats :: [Int]
nats = 1 : map (+1) nats

Those two examples are both referring to earlier parts of the list, but it's also possible to refer to a later part of the list.

-- |
-- >>> backwards
-- [7,8,9,10]
backwards :: [Int]
backwards = [ (backwards !! 1) - 1
            , (backwards !! 2) - 1
            , (backwards !! 3) - 1
            , 10
            ]

Here is what happens while backwards is being printed. First, ((backwards !! 1) - 1)'s thunk is forced. This causes (backwards !! 2) - 1 to be forced, which in turn causes (backwards !! 3) - 1 to be forced as well. This yields the result 9, so the (backwards !! 3) - 1 thunk is replaced with the value 9, and similarly for 8 and 7. The first element 7 can finally be printed, and we move on to print the next elements, which have already been evaluated to 8, 9 and 10. As you can see, each thunk is only evaluated once, and if there are dependencies between the thunks, the dependencies are evaluated first.

I hope it is now clear how laziness can help us with today's puzzle: instead of writing an imperative computation to manually cache the results and traverse the dependencies, we can simply construct a data structure whose elements are defined in terms of each other. This way, evaluating any element of the data structure will automatically force its dependencies while caching the results.

day7 :: [Assignment] -> WireName -> Word16
day7 assignments name0 = lookupCache' cache name0
  where
    cache :: Cache
    cache = Map.fromList (map (second eval) assignments)
    
    eval :: Expr -> Word16
    eval (Number   n) = n
    eval (Wire  name) = lookupCache' cache name
    eval (AND    x y) = opAND    (eval x) (eval y)
    eval (OR     x y) = opOR     (eval x) (eval y)
    eval (LSHIFT x i) = opLSHIFT (eval x) i
    eval (RSHIFT x i) = opRSHIFT (eval x) i
    eval (NOT    x  ) = opNOT    (eval x)

Here I am constructing cache, a map from each wire name to its computed value, by applying eval to the corresponding expressions. eval is itself defined in terms of the cache, in which it looks up the values computed for the dependencies. We didn't have to write any special code to do so, but this cache lookup will either return immediately with the result if it has already been computed, or it will find and force a thunk, calling eval on the dependency and caching the result.

Another way to achieve the same result is to use a memoization library. Here I use MemoTrie to memoize the result of looking up and evaluating the expression corresponding to a wire name. This way, when evalExpr calls evalWire on a dependency, it will use the cached value if there is one. The implementation uses laziness under the hood.

import Data.MemoTrie

day7 :: [Assignment] -> WireName -> Word16
day7 assignments name0 = evalWire name0
  where
    env :: Env
    env = mkEnv assignments
    
    evalWire :: WireName -> Word16
    evalWire = memo (evalExpr . lookupEnv' env)
    
    evalExpr :: Expr -> Word16
    evalExpr (Number   n) = n
    evalExpr (Wire  name) = evalWire name
    evalExpr (AND    x y) = opAND    (evalExpr x) (evalExpr y)
    evalExpr (OR     x y) = opOR     (evalExpr x) (evalExpr y)
    evalExpr (LSHIFT x i) = opLSHIFT (evalExpr x) i
    evalExpr (RSHIFT x i) = opRSHIFT (evalExpr x) i
    evalExpr (NOT    x  ) = opNOT    (evalExpr x)

If you'd like to read more about using laziness for caching, I have written a post explaining when caches do and do not get reused between calls.

Day 7, meta-programmatically: Template Haskell

I know this post is getting quite long already, but I'd like to point out one last way in which the puzzle could be solved. Consider the following Haskell code:

d = opAND x y
e = opOR x y
f = opLSHIFT x 2
g = opRSHIFT y 2
h = opNOT x
i = opNOT y
x = 123
y = 456

Even though some of the above values are used before being defined, the thunk forcing mechanism makes sure that evaluating any of those values will first evaluate and cache its dependencies. So if we could generate similar Haskell code from the circuit description, we could efficiently obtain the value of any wire.

One way of generating code which is available in any language is to construct a string containing the code we want:

import Text.Printf

validName :: String -> String
validName "do" = "do'"
validName "id" = "id'"
validName "if" = "if'"
validName "in" = "in'"
validName name = name

exprS :: Expr -> String
exprS (Number   n) = show n
exprS (Wire  name) = validName name
exprS (AND    x y) = printf "opAND    %s %s" (exprS x) (exprS y)
exprS (OR     x y) = printf "opOR     %s %s" (exprS x) (exprS y)
exprS (LSHIFT x i) = printf "opLSHIFT %s %d" (exprS x) i
exprS (RSHIFT x i) = printf "opRSHIFT %s %d" (exprS x) i
exprS (NOT    x  ) = printf "opNOT    %s"    (exprS x)

-- |
-- >>> assignmentS ("g", RSHIFT (Wire "y") 2)
-- "g = opRSHIFT y 2"
assignmentS :: Assignment -> String
assignmentS (name, e) = printf "%s = %s" name (exprS e)

This works, but generating strings is quite error-prone, in the sense that the compiler is not going to prevent you from generating code which doesn't compile.

A more principled approach is to use Template Haskell. Instead of manipulating strings, we manipulate quasiquotations which fail at compile time if the quoted code does not parse.

{-# LANGUAGE TemplateHaskell #-}
import Language.Haskell.TH

validName :: String -> Name
validName "do" = mkName "do'"
validName "id" = mkName "id'"
validName "if" = mkName "if'"
validName "in" = mkName "in'"
validName name = mkName name

exprE :: Expr -> Q Exp
exprE (Number   n) = [| n |]
exprE (Wire  name) = varE (validName name)
exprE (AND    x y) = [| opAND    $(exprE x) $(exprE y) |]
exprE (OR     x y) = [| opOR     $(exprE x) $(exprE y) |]
exprE (LSHIFT x i) = [| opLSHIFT $(exprE x) i          |]
exprE (RSHIFT x i) = [| opRSHIFT $(exprE x) i          |]
exprE (NOT    x  ) = [| opNOT    $(exprE x)            |]

-- >>> let assignment = ("g", RSHIFT (Wire "y") 2)
-- >>> decls <- runQ $ assignmentD assignment
-- >>> ppr decls
-- g = opRSHIFT y 2
assignmentD :: Assignment -> Q [Dec]
assignmentD (name, e) = [d| $(var) = $(val) |]
  where
    var :: Q Pat
    var = varP (validName name)
    
    val :: Q Exp
    val = exprE e

For extra safety, there is also a variant of [| ... |] which detects type errors at compile time:

import Language.Haskell.TH.Syntax

exprE :: Expr -> Q (TExp Word16)
exprE (Number   n) = [|| n ||]
exprE (Wire  name) = unsafeTExpCoerce $ varE (validName name)
exprE (AND    x y) = [|| opAND    $$(exprE x) $$(exprE y) ||]
exprE (OR     x y) = [|| opOR     $$(exprE x) $$(exprE y) ||]
exprE (LSHIFT x i) = [|| opLSHIFT $$(exprE x) i           ||]
exprE (RSHIFT x i) = [|| opRSHIFT $$(exprE x) i           ||]
exprE (NOT    x  ) = [|| opNOT    $$(exprE x)             ||]

Since we're generating variable names dynamically from the wire names, Template Haskell cannot know at compile time whether those variables have the correct type, so we have to use unsafeTExpCoerce to tell it to trust us that this piece will work out.

To complete the puzzle, I write a small wrapper which loads the circuit description from a file and concatenates all the generated declarations:

circuit :: FilePath -> Q [Dec]
circuit filename = do
    Right assignments <- runIO
                       $ parseFromFile (many assignment) filename
    decls <- forM assignments assignmentD
    return (concat decls)

I can finally load this circuit into my main program (for technical reason, the definitions above must be in a different file from the code below) and access the variables it defines.

-- defines "a" and many other variables
$(circuit "input.txt")

day7 :: Word16
day7 = a

Navigation

• About this series
• Previous day: precise types, arrays, and in-place updates
• Next day: invertible parsers