Writing an assembler in Haskell - mapM with state? - haskell

Question

I'm writing a very simple two-pass assembler in Haskell and I've come across a scenario that I don't yet have the experience to solve. I think the solution is likely to involve monad transformers, which I don't really understand.
The assembler parses the assembly code into a list of Statements, which are either instructions or labels. Some Statements may refer to labels. The assembler needs to convert the Statements into Instructions, which involves eliminating the labels and substituting the label references with an appropriate value.
I have written the first pass of the assembler, which produces a [(String, Int)] representing a map from labels to addresses. I have also written the following function for translating a Statement into an Instruction:
stmtToInstruction :: Int -> [(String, Int)] -> Statement -> Either String [I.Instruction]
stmtToInstruction addr labels stmt = case stmt of
ADD d s1 s2 -> Right [I.ADD d s1 s2]
BEQL s1 s2 l -> case do label <- find (\e -> fst e == l) labels
let labelAddr = snd label
let relativeAddr = I.ImmS $ fromIntegral (labelAddr - addr)
return (I.BEQ s1 s2 relativeAddr) of
Just i -> Right [i]
Nothing -> Left $ "Label " ++ l ++ " not defined"
LABEL _ -> Right []
I've omitted several cases for brevity, but you can see all the possible results here:
ADD always succeeds and produces an instruction
BEQL can either succeed or fail, depending on whether a label is found
LABEL always succeeds, even though it produces no actual instructions
This works as expected. The problem I now have is writing this function:
replaceLabels :: [Statement] -> Either String [I.Instruction]
replaceLabels takes a list of statements, and runs stmtToInstruction on each one. The addr argument to stmtToInstruction must be the length of the [Instruction] accumulated so far. The output may either be a Left String, if one of the label references was invalid, or a Right [I.Instruction], if there were no errors.
mapM :: Monad m => (a -> m b) -> [a] -> m [b] gets us some of the way there, but provides no way to inject the current address into the (a -> m b) function. How do I make this work?

Answer 1

You're right: the StateT monad transformer will do the trick:
imapM :: (Traversable t, Monad m)
=> (Int -> a -> m b) -> t a -> m (t b)
imapM f = flip runStateT 0 .
mapM (\a ->
do
count <- get
put $! count + 1
f count a)
But writing the specialized version for lists might be better:
itraverse :: Applicative f
=> (Int -> a -> f b) -> [a] -> f [b]
itraverse f = go 0 where
go !_ [] = pure []
go !count (x:xs) = (:) <$> f count x <*> go (count + 1) xs

Answer 2

I've implemented a recursive solution that I'm sure is very inefficient. I'd still be interested to see the 'proper' way of doing this.
replaceLabels :: [Statement] -> Either String [I.Instruction]
replaceLabels [] = Right []
replaceLabels stmts#(s:ss) = replaceLabels' labels stmts 0
where labels = process stmts
replaceLabels' :: [(String, Int)] -> [Statement] -> Int -> Either String [I.Instruction]
replaceLabels' _ [] _ = Right []
replaceLabels' labels (s:ss) addr = do
instructions <- stmtToInstruction addr labels s
restInstructions <- replaceLabels' labels ss (addr + length instructions)
return (instructions ++ restInstructions)

Answer 3

I would start by changing
stmtToInstruction :: Int -> [(String, Int)] -> Statement -> Either String [I.Instruction]
into
stmtToInstruction :: [(String, Int)] -> Statement -> Either String (Int -> [I.Instruction])
That is, moving the function that takes the address into the Right branch of the Either. The reason is that label reference errors seem to be independent of addresses, so it's better to handle reference errors first and then worry about the address stuff in isolation.
This function resolves the references:
resolveRefs :: [(String,Int)] -> [Statement] -> Either String [Int -> [Instruction]]
resolveRefs environment = traverse (stmtToInstruction environment)
(traverse is equivalent to mapM but it only requires an Applicative constraint. They are different functions merely for historical reasons.)
Ok, after having handled the errors, lets now focus on the [Int -> [Instruction]] list. It seems that we have to map over it from the left while carrying an accumulated address that we must supply to each function. The mapAccumL function is perfect for this:
resolveAddrs :: [Int -> [Instruction]] -> [Instruction]
resolveAddrs funcs = mconcat . snd $ accumulate funcs
where
accumulate :: [Int -> [Instruction]] -> (Int,[[Instruction]])
accumulate = mapAccumL step 0
step address func = let is = func address in (address + length is,is)