assembly_demo.hs

import Data.Bits
import Data.List
import Data.List.Split
import qualified Data.Map as Map
import Data.Maybe

-- Abstract operation type
data Operation = Operation {
    compute :: [Integer] -> Integer,
    symbol :: String
}
instance Show Operation where
    show = symbol
instance Eq Operation where
    x == y = symbol x == symbol y

-- Instruction abstract type
    -- take data from `inputs` operands
    -- perform an `operation` on it
    -- put the result in `output` register
    -- this is a non-recursive definition
    -- instructions only have meaning in context with other instructions
data Instruction = Instruction {
    inputs :: [Operand],
    output :: Register,
    operation :: Operation
} deriving Show

-- Operand abstract type
    -- An operand can be an integer immediate value or a register reference
data Operand = Operand Integer | Register Register
    deriving (Show)

-- Register abstract type
    -- At the moment, just a wrapper for Integer
data Register = R Integer
    deriving (Eq, Ord, Show)

-- Computation abstract type
    -- perform an operation on the results of subcomputations
    -- this is a recursive definition with a tree structure
    -- computations are abstract and have meaning without context
data Computation = Data Integer | Var Integer | Exp {
    dependencies :: [Computation],
    operator :: Operation
}
instance Show Computation where
    show (Data d) = show d
    show (Var x) = "V" ++ show x
    show (Exp deps op) = (show deps) ++ (show op)

-- Abstract memory type
    -- Maps registers to computations describing their values
type AbstractMem = (Map.Map Register Computation)

-- Definition of add operation and instruction
addOp = Operation {
    compute = \[x, y] -> x + y,
    symbol = "+"
}
makeAdd :: Register -> Operand -> Operand -> Instruction
makeAdd out in1 in2 = Instruction {
    inputs = [in1, in2],
    output = out,
    operation = addOp
}

-- Definition of inv operation and instruction
invOp = Operation {
    compute = \[x] -> -x,
    symbol = "-"
}
makeInv :: Register -> Operand -> Instruction
makeInv out in1 = Instruction {
    inputs = [in1],
    output = out,
    operation = invOp
}

-- Definition of and operation and instruction
andOp = Operation {
    compute = \[x, y] -> x .&. y,
    symbol = "&"
}
makeAnd :: Register -> Operand -> Operand -> Instruction
makeAnd out in1 in2 = Instruction {
    inputs = [in1, in2],
    output = out,
    operation = andOp
}

-- Definition of not operation and instruction
notOp = Operation {
    compute = \[x] -> complement x,
    symbol = "~"
}
makeNot :: Register -> Operand -> Instruction
makeNot out in1 = Instruction {
    inputs = [in1],
    output = out,
    operation = notOp
}

-- Break program into lines and words and readLine
readProgram :: String -> [Instruction]
readProgram s = map (readLine . (splitOn " ")) $ splitOn "\n" s

-- Take a list of words in a line and parse it to the appropriate instruction
readLine :: [String] -> Instruction
readLine ("add":out:in1:in2:[]) = makeAdd (readReg out) (readOp in1) (readOp in2)
readLine ("inv":out:in1:[]) = makeInv (readReg out) (readOp in1)
readLine ("and":out:in1:in2:[]) = makeAnd (readReg out) (readOp in1) (readOp in2)
readLine ("not":out:in1:[]) = makeNot (readReg out) (readOp in1)

-- Parse word to operand
readOp :: String -> Operand
readOp ('$':x) = Register $ readReg ('$':x)
readOp x = Operand $ read x

-- Parse word to register
readReg :: String -> Register
readReg ('$':x) = R $ read x

-- Transform a list of instructions to an abstract memory map
    -- Basically AbstractMem is a data structure representing
    --  the program's memory at a certain point in time
    -- We can iterate through each instruction, updating the
    --  memory with computations as we go
    -- With each step, the computations build on each other,
    --  creating an increasingly nested tree structure
abstract :: [Instruction] -> AbstractMem -> AbstractMem
abstract (instruction:instructions) mem = abstract instructions nextMem
    where nextMem = Map.insert (output instruction) computation mem
          computation = Exp {
              operator = (operation instruction),
              dependencies = map (toComputation mem) $ inputs instruction
          }
abstract [] mem = mem

-- Dereference registers and operands
--  Undefined registers are treated as variables
toComputation :: AbstractMem -> Operand -> Computation
toComputation mem (Register (R x)) = fromMaybe (Var x) $ Map.lookup (R x) mem
toComputation mem (Operand x) = Data x

-- Is the computation just a number?
isData :: Computation -> Bool
isData (Data _) = True
isData _ = False

-- If a computation has no variables, we can just compute the result now
    -- Otherwise, try to collapse the subtrees the same way
evaluate :: Computation -> Computation
evaluate (Exp deps op)
    | all isData evDeps = Data $ (compute op) $ map (\(Data x) -> x) evDeps
    | otherwise = Exp evDeps op
    where evDeps = map evaluate deps
evaluate x = x

-- Takes expressions of the form ((x + a) + b) and transforms them to ((a + b) + x)
    -- Where x is definitely a variable or expression, and a and b might not be
    -- Also works for &
    -- This tends to move variables higher in the tree, so variable-free collapsible
    --  subtrees become bigger and we can simplify more
associate :: Computation -> Computation
associate (Exp [Exp [x@(Var _), a] op1, b] op2) | (op1 == op2 && (op1 == addOp || op1 == andOp))
    = Exp [x, Exp [a, b] op1] op2
associate (Exp [Exp [x@(Exp _ _), a] op1, b] op2) | (op1 == op2 && (op1 == addOp || op1 == andOp))
    = Exp [x, Exp [a, b] op1] op2
associate (Exp deps op) = (Exp (map associate deps) op)
associate x = x

-- Takes expressions of the form (a + x) and transforms them to (x + a)
    -- Where x is definitely a variable or expression, and a might not be
    -- This tends to standardize the form so more complex branches come first
    -- This helps associate do its job better
commute :: Computation -> Computation
commute (Exp [a, x@(Var _)] op) | (op == addOp || op == andOp)
    = Exp [x, a] op
commute (Exp [a, x@(Exp _ _)] op) | (op == addOp || op == andOp)
    = Exp [x, a] op
commute (Exp deps op) = (Exp (map commute deps) op)
commute x = x


main = do
    let prog = "add $7 $8 1\n\
               \not $5 0\n\
               \inv $6 4\n\
               \add $7 $7 $6\n\
               \and $6 $6 $5\n\
               \inv $6 $6\n\
               \add $7 $7 $9\n\
               \add $7 $7 $6"
    -- Abstract the program and get the computation for register 7:
    let r7 = toComputation (abstract (readProgram prog) Map.empty) (Register $ R 7)
    -- Iteratively optimize it
    let optimizedR7 = iterate (evaluate . commute . associate . commute) r7
    -- Print the first expression after each 0-4 passes of the optimizer
    putStrLn $ concat $ intersperse "\n" $ map show $ take 5 $ optimizedR7