Simple register allocation, see "Essentials of Compilation" for more https://jeapostrophe.github.io/courses/2017/spring/406/notes/book.pdf

coloring-3reg.svg

<?xml version="1.0" encoding="UTF-8" standalone="no"?>
<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"
 "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
<!-- Generated by graphviz version 2.38.0 (20140413.2041)
 -->
<!-- Title: %3 Pages: 1 -->
<svg width="142pt" height="476pt"
 viewBox="0.00 0.00 141.95 476.00" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
<g id="graph0" class="graph" transform="scale(1 1) rotate(0) translate(4 472)">
<title>%3</title>
<polygon fill="white" stroke="none" points="-4,4 -4,-472 137.946,-472 137.946,4 -4,4"/>
<!-- 1 -->
<g id="node1" class="node"><title>1</title>
<ellipse fill="none" stroke="#cc5151" cx="51.9464" cy="-18" rx="27" ry="18"/>
<text text-anchor="middle" x="51.9464" y="-14.3" font-family="Times,serif" font-size="14.00">v</text>
</g>
<!-- 2 -->
<g id="node2" class="node"><title>2</title>
<ellipse fill="none" stroke="black" stroke-dasharray="5,2" cx="51.9464" cy="-90" rx="43.5923" ry="18"/>
<text text-anchor="middle" x="51.9464" y="-86.3" font-family="Times,serif" font-size="14.00">w (stack)</text>
</g>
<!-- 2&#45;&gt;1 -->
<g id="edge1" class="edge"><title>2&#45;&gt;1</title>
<path fill="none" stroke="black" d="M51.9464,-71.6966C51.9464,-60.8463 51.9464,-46.9167 51.9464,-36.1043"/>
</g>
<!-- 3 -->
<g id="node3" class="node"><title>3</title>
<ellipse fill="none" stroke="#cc5151" cx="106.946" cy="-162" rx="27" ry="18"/>
<text text-anchor="middle" x="106.946" y="-158.3" font-family="Times,serif" font-size="14.00">x</text>
</g>
<!-- 3&#45;&gt;2 -->
<g id="edge2" class="edge"><title>3&#45;&gt;2</title>
<path fill="none" stroke="black" d="M95.0122,-145.811C86.166,-134.552 74.1155,-119.215 65.0111,-107.628"/>
</g>
<!-- 4 -->
<g id="node4" class="node"><title>4</title>
<ellipse fill="none" stroke="#7f3333" cx="68.9464" cy="-234" rx="27" ry="18"/>
<text text-anchor="middle" x="68.9464" y="-230.3" font-family="Times,serif" font-size="14.00">y</text>
</g>
<!-- 4&#45;&gt;2 -->
<g id="edge3" class="edge"><title>4&#45;&gt;2</title>
<path fill="none" stroke="black" d="M66.8958,-215.871C63.6283,-188.578 57.2763,-135.52 54.0042,-108.189"/>
</g>
<!-- 4&#45;&gt;3 -->
<g id="edge4" class="edge"><title>4&#45;&gt;3</title>
<path fill="none" stroke="black" d="M77.7598,-216.765C83.8988,-205.456 92.0522,-190.437 98.1807,-179.147"/>
</g>
<!-- 5 -->
<g id="node5" class="node"><title>5</title>
<ellipse fill="none" stroke="black" stroke-dasharray="5,2" cx="40.9464" cy="-306" rx="40.8928" ry="18"/>
<text text-anchor="middle" x="40.9464" y="-302.3" font-family="Times,serif" font-size="14.00">z (stack)</text>
</g>
<!-- 5&#45;&gt;2 -->
<g id="edge5" class="edge"><title>5&#45;&gt;2</title>
<path fill="none" stroke="black" d="M37.698,-287.888C35.9005,-277.542 33.8517,-264.063 32.9464,-252 29.0076,-199.516 40.9612,-137.801 47.6911,-108.345"/>
</g>
<!-- 5&#45;&gt;4 -->
<g id="edge6" class="edge"><title>5&#45;&gt;4</title>
<path fill="none" stroke="black" d="M47.7244,-288.055C52.1267,-277.049 57.8409,-262.764 62.2308,-251.789"/>
</g>
<!-- 6 -->
<g id="node6" class="node"><title>6</title>
<ellipse fill="none" stroke="#cc5151" cx="40.9464" cy="-378" rx="27" ry="18"/>
<text text-anchor="middle" x="40.9464" y="-374.3" font-family="Times,serif" font-size="14.00">t1</text>
</g>
<!-- 6&#45;&gt;5 -->
<g id="edge7" class="edge"><title>6&#45;&gt;5</title>
<path fill="none" stroke="black" d="M40.9464,-359.697C40.9464,-348.846 40.9464,-334.917 40.9464,-324.104"/>
</g>
<!-- 7 -->
<g id="node7" class="node"><title>7</title>
<ellipse fill="none" stroke="#7f3333" cx="40.9464" cy="-450" rx="27" ry="18"/>
<text text-anchor="middle" x="40.9464" y="-446.3" font-family="Times,serif" font-size="14.00">t2</text>
</g>
<!-- 7&#45;&gt;6 -->
<g id="edge8" class="edge"><title>7&#45;&gt;6</title>
<path fill="none" stroke="black" d="M40.9464,-431.697C40.9464,-420.846 40.9464,-406.917 40.9464,-396.104"/>
</g>
</g>
</svg>

example analysis

aseipp@ubuntu:~/t$ ghci RegAlloc1.hs 
GHCi, version 8.0.2: http://www.haskell.org/ghc/  :? for help
[1 of 1] Compiling RegAlloc1        ( RegAlloc1.hs, interpreted )
Ok, modules loaded: RegAlloc1.
*RegAlloc1> printLive program1 
    movq $1, v.1    | {v}
    movq $46, w.2   | {v,w}
    movq v.1, x.3   | {w,x}
    addq $7, x.3    | {w,x}
    movq x.3, y.4   | {w,x,y}
    addq $4, y.4    | {w,x,y}
    movq x.3, z.5   | {w,y,z}
    addq w.2, z.5   | {y,z}
    movq y.4, t1.6  | {z,t1}
    negq t1.6       | {z,t1}
    movq z.5, t2.7  | {t1,t2}
    addq t1.6, t2.7 | {t2}
    movq t2.7, %rax | {}
*RegAlloc1> let (l, i, c) = (liveness program1, interference program1 l, colorGraph 3 i) in writeDotColored "coloring.dot" c
*RegAlloc1> 

RegAlloc1.hs

{-# OPTIONS_GHC -Wall #-}

{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE PartialTypeSignatures #-}
module RegAlloc1
  ( -- * Types
    Var
  , Reg(..)
  , Arg(..)
  , Instr(..)
  , Program(..)

    -- * Liveness analysis
  , Live
  , liveness
  , printLive

    -- * Interference
  , Interference
  , interference

    -- * Graph coloring
  , colorGraph

    -- * Example programs
  , program1

    -- * Utilities
    -- ** GraphViz
  , writeDot
  , writeDotInterference
  , writeDotColored

    -- ** x86 Registers
  , rax, rdi, rsi, rdx, rcx, r8, r9
  ) where

import           Control.Monad

import           Data.Data
import           Data.Function
import           Data.List
import           Data.Monoid
import           Data.String

-- containers
import           Data.Map            ( Map )
import qualified Data.Map            as Map
import           Data.Set            ( Set )
import qualified Data.Set            as Set

-- fgl, graphviz
import           Data.Graph.Inductive.Graph
import           Data.Graph.Inductive.PatriciaTree
import qualified Data.GraphViz                       as GV
import qualified Data.GraphViz.Attributes.Complete   as GV
import qualified Data.GraphViz.Commands.IO           as GV

--------------------------------------------------------------------------------
-- basics

data Var = VarU { _varI :: !Int, _varS :: String }
  deriving (Eq, Ord, Data, Typeable)

instance Show Var where
  show (VarU i s) = s <> ('.':show i)

data Reg = Rax | Rdi | Rsi | Rdx | Rcx | R8 | R9
  deriving (Eq, Ord, Data, Typeable)

instance Show Reg where
  show Rax = "%rax"
  show Rdi = "%rdi"
  show Rsi = "%rsi"
  show Rdx = "%rdx"
  show Rcx = "%rcx"
  show R8  = "%r8"
  show R9  = "%r9"

data Arg = LitI Integer | Var Var | Reg Reg
  deriving (Eq, Ord, Data, Typeable)

instance Show Arg where
  show (LitI i) = '$':show i
  show (Var v)  = show v
  show (Reg r)  = show r

data Instr
  = Movq Arg Arg
  | Addq Arg Arg
  | Negq Arg
  deriving (Eq, Ord, Data, Typeable)

instance Show Instr where
  show (Movq a d) = "movq" <+> show a <> "," <+> show d
  show (Addq a d) = "addq" <+> show a <> "," <+> show d
  show (Negq a)   = "negq" <+> show a

data Program = Program String [Instr]
  deriving (Eq, Ord, Data, Typeable)

instance Show Program where
  show (Program name instrs) =
    let fullName = '_':name
        indent   = ("    "++)

    in   ".global" <+> fullName
    <++> fullName <> ":"
    <++> unlines (map (indent . show) instrs)

--
-- utilities
--

rax, rdi, rsi, rdx, rcx, r8, r9 :: Arg
(rax, rdi, rsi, rdx, rcx, r8, r9)
  = (Reg Rax, Reg Rdi, Reg Rsi, Reg Rdx, Reg Rcx, Reg R8, Reg R9)

(<+>) :: (Monoid m, IsString m) => m -> m -> m
(<+>) a b  = a <> " "  <> b

(<++>) :: (Monoid m, IsString m) => m -> m -> m
(<++>) a b = a <> "\n" <> b

--
-- useful but stupid instances for writing clear examples
--

instance Num Arg where
  fromInteger = LitI

  -- warnings, lol
  (+)    = error "lol"
  (*)    = error "lol"
  abs    = error "lol"
  signum = error "lol"
  negate = error "lol"

--------------------------------------------------------------------------------
-- example program

program1 :: Program
program1 = Program "prog1"
  [ Movq 1  v
  , Movq 46 w
  , Movq v  x
  , Addq 7  x
  , Movq x  y
  , Addq 4  y
  , Movq x  z
  , Addq w  z
  , Movq y  t1
  , Negq t1
  , Movq z  t2
  , Addq t1 t2
  , Movq t2 rax
  ]
  where
    v  = Var $ VarU 1 "v"
    w  = Var $ VarU 2 "w"
    x  = Var $ VarU 3 "x"
    y  = Var $ VarU 4 "y"
    z  = Var $ VarU 5 "z"
    t1 = Var $ VarU 6 "t1"
    t2 = Var $ VarU 7 "t2"

--------------------------------------------------------------------------------
-- liveness analysis

type Live = Set Var

data LiveSet = Read | Write

getVar :: Arg -> Set Var
getVar (Var a) = Set.singleton a
getVar _       = Set.empty

need :: LiveSet -> Instr -> Live
-- the variables that are read in an instruction
need Read (Movq a _) = getVar a
need Read (Addq a d) = getVar a `Set.union` getVar d
need Read (Negq a)   = getVar a

-- the variables that are written in an instruction
need Write (Movq _ d) = getVar d
need Write (Addq _ d) = getVar d
need Write (Negq a)   = getVar a

-- | Perform a liveness analysis on the input program and determine the set of
-- live variables for every point in the program.
liveness :: Program -> [Live]
liveness (Program _ instrs)
  = reverse $ Set.empty : go input (k end Set.empty)
  where
    (end:input) = reverse instrs

    k x after = (after `Set.difference` w) `Set.union` r
      where (r, w) = (need Read x, need Write x)

    go []     _     = []
    go (x:xs) after = after : go xs (k x after)

-- | Print a program and its liveness information. Useful for debugging.
printLive :: Program -> IO ()
printLive prog@(Program _ instrs)
  = void
  $ flip traverse (zip instrs $ liveness prog)
  $ \(instr, lv) -> do
    let msg    = "    " ++ show instr
        buffer = replicate (20 - length msg) ' '
        info   = intercalate "," $ map _varS (Set.toList lv)
    putStrLn (msg ++ buffer ++ "| {" ++ info ++ "}")

--------------------------------------------------------------------------------
-- interference graphs

-- | An interference graph for a given @'Program'@. Nodes in the graph (which
-- are @Int@s) are labeled with the original variable they represented. Edges
-- are undirected, and have no labels as they convey no meaning.
type Interference = Gr Var ()

-- | Build an interference graph for a given program, based on the results of a
-- liveness analysis. Two variables in the input program interfere -- they are
-- both live at the same time -- if those variables have an edge between them in
-- the resulting graph.
interference :: Program
             -- ^ the input program
             -> [Live]
             -- ^ results of a liveness analysis on the input program
             -> Interference
             -- ^ interference graph
interference (Program _ instrs) live = foldr loop empty (zip instrs live)
  where
    -- create an edge between two variables. does not update the graph if the
    -- edge already exists.
    mkIEdge (VarU a _) (VarU b _) gr
      | hasEdge gr (a, b) = gr
      | otherwise         = insEdge (a, b, ()) gr

    -- insert a node if it doesn't exist already. does not update the graph if
    -- the given node already exists.
    mkINode v@(VarU i _) gr
      | gelem i gr = gr
      | otherwise  = insNode (i, v) gr

    -- insert a list of nodes, all at once.
    mkINodes vs gr = foldr mkINode gr vs

    -- mark two variables as interfering with each other, e.g. they are both
    -- live at the same time. this is represented as a single edge between both
    -- nodes in the graph.
    interferes :: Var -> Var -> Interference -> Interference
    interferes v1 v2 gr
      = mkIEdge v1 v2     -- create an edge between them
      $ mkINodes [v1, v2] -- insert both nodes
      $ gr

    -- mark a variable @d@ as interfering with all of the variables @v@ in given
    -- live variable set @lv@, iff @v@ passes a given predicate. this
    -- essentially creates a one-to-many mapping between @d@ and the elements of
    -- set @lv@ in the graph.
    insertP :: (Var -> Bool)
            -- ^ predicate function. for every variable @v \elem vs@ in
            -- the live set @vs@, if @v@ passes this predicate, then the two
            -- variables @(d, v)@ are marked as interfering, i.e. an edge is
            -- created between them.
            -> (Var, Live)
            -- ^ @(d, lv)@ -- the variable @d@ and live variable set @lv@
            -> Interference
            -- ^ input graph
            -> Interference
            -- ^ resulting graph
    insertP predicate (d, lv) gr =
        let k :: Var -> Interference -> Interference
            k v g | predicate v = interferes d v g
                  | otherwise   = g
        in Set.foldr k gr lv

    -- mark all of the interferences for a given instruction and its live
    -- variable set.
    loop :: (Instr, Live)
         -- ^ input instruction, and live variable set for that instruction
         -> Interference
         -- ^ input graph
         -> Interference
         -- ^ output graph
    loop x gr = case x of
      -- for move instructions, we mark @d@ as interfering with @v \elem lv@,
      -- iff @v /= d@ (we don't associate @d@ with itself) and @d /= a@ (we
      -- don't associate @d@ with the source @a@, since @a@ may be dead after
      -- this, or it may not be a variable at all)
      (Movq a (Var d), lv) -> insertP (\v -> v /= d && (Var v) /= a) (d, lv) gr

      -- for arithmetic instructions, we mark @d@ as interfering with
      -- @v \elem lv@, iff @v /= d@ (we don't associate @d@ with itself)
      (Addq _ (Var d), lv) -> insertP (\v -> v /= d)                 (d, lv) gr
      (Negq   (Var d), lv) -> insertP (\v -> v /= d)                 (d, lv) gr

      -- otherwise, return
      _ -> gr

writeDotInterference :: FilePath -> Interference -> IO ()
writeDotInterference = writeDot $ \l -> [ GV.toLabel (_varS l) ]

--------------------------------------------------------------------------------
-- Graph coloring via DSATUR

type GColor  = (Int, GV.Color)
type Colored = Gr (Var, Maybe GColor) ()

-- see https://stackoverflow.com/a/13781114
allColors :: [GV.Color]
allColors = do
  i <- [ 2** k | k <- [0..] ]
  j <- enumFromThenTo 1 3 i
  v <- [ 8 / 10, 5 / 10 ]
  return $ GV.HSV (j / i) (6 / 10) v

color :: Int -> Node -> GColor
color limit i = zip [0..] allColors !! (i `mod` limit)

saturation :: Graph gr => Int -> gr a b -> Node -> [Node]
saturation limit gr u
  = map fst $ ordNubBy id ((==) `on` fst)
      [ color limit v
      | v <- neighbors gr u
      ]

sortForColoring :: Graph gr => Int -> gr a b -> [Node]
sortForColoring limit gr
  = sortBy (flip compare `on` length . saturation limit gr)
  $ nodes gr

getMinColor :: Graph gr => Int -> gr a b -> Node -> Maybe GColor
getMinColor limit gr u
  = let taken      = saturation limit gr u -- find all the colors that are taken
        everyColor = [ 0 .. limit-1 ]      -- set of all possible colors
        getColor   = color limit
    in case everyColor \\ taken of          -- find non-taken colors
         [] -> Nothing                      -- none available; stack slot
         vs -> Just (getColor $ minimum vs) -- otherwise, get min color

buildColorMap :: Int -> Interference -> Map Node (Maybe GColor)
buildColorMap limit gr
  = foldr (\n -> Map.insert n (getMinColor limit gr n)) Map.empty
  $ sortForColoring limit gr

colorGraph :: Int -> Interference -> Colored
colorGraph limit gr = nmap (\v@(VarU i _) -> (v, get i)) gr
  where
    get       = maybe (error "colorGraph: impossible!") id . findMap
    findMap x = Map.lookup x (buildColorMap limit gr)

writeDotColored :: FilePath -> Colored -> IO ()
writeDotColored = writeDot $ \(l, c) ->
  case c of
    -- no coloring; must be a stack slot
    Nothing      -> [ GV.toLabel $ _varS l <+> "(stack)"
                    , GV.style GV.dashed
                    ]
    Just (_, cl) -> [ GV.toLabel $ _varS l
                    , GV.Color $ GV.toColorList [ cl ]
                    , GV.style GV.solid
                    ]

--------------------------------------------------------------------------------
-- GraphViz utilities

writeDot :: Graph gr
         => (a -> GV.Attributes)
         -> FilePath
         -> gr a b
         -> IO ()
writeDot markup fp
  = GV.writeDotFile fp
  . GV.graphToDot GV.nonClusteredParams
      { GV.globalAttributes = [ noDir ]   -- undirected edges
      , GV.fmtEdge = \(_, _, _l) -> []    -- edges have no label
      , GV.fmtNode = \(_, l) -> markup l  -- label nodes by name
      }
  where
    noDir = GV.EdgeAttrs [ GV.Dir GV.NoDir ]

ordNubBy :: (Ord b) => (a -> b) -> (a -> a -> Bool) -> [a] -> [a]
ordNubBy p f l = go Map.empty l
  where
    go _ []     = []
    go m (x:xs) =
      let b = p x
      in case b `Map.lookup` m of
           Nothing     -> x : go (Map.insert b [x] m) xs
           Just bucket
             | elem_by f x bucket -> go m xs
             | otherwise          -> x : go (Map.insert b (x:bucket) m) xs

    -- From the Data.List source code.
    elem_by :: (a -> a -> Bool) -> a -> [a] -> Bool
    elem_by _  _ []     = False
    elem_by eq y (x:xs) = y `eq` x || elem_by eq y xs