1 files changed, 87 insertions, 0 deletions
diff --git a/RegAlloc.hs b/RegAlloc.hs
new file mode 100644
index 0000000..d2b1717
--- /dev/null
+++ b/RegAlloc.hs
@@ -0,0 +1,87 @@
+module RegAlloc(regalloc, Allocation(..)) where
+
+import Data.Function
+import Data.List
+import Data.Maybe
+import qualified Data.Map.Strict as Map
+import Debug.Trace
+
+import Utils
+
+
+-- Follows the Linear Scan Register Allocation algorithm specified in:
+-- https://www.cs.purdue.edu/homes/suresh/502-Fall2008/papers/linear-scan.pdf
+
+-- Test case:
+-- regalloc [((1, 3), "a"), ((2, 2), "x"), ((2, 10), "b"), ((3, 7), "c"), ((4, 5), "d")] ["ra", "rb"]
+
+
+data Allocation a = AllocReg a | AllocMem
+  deriving Show
+
+type Interval = (Int, Int)
+
+data State a = State {stActive :: [Int], stAlloc :: [Allocation a], stFreeRegs :: [a]}
+  deriving Show
+
+regalloc :: (Show a, Show b, Ord a, Ord b)
+         => [(Interval, b)]           -- [(live interval, name of variable)]
+         -> [a]                       -- the available registers
+         -> [(b, b)]                  -- pairs to be allocated to the same register if possible
+         -> Map.Map b (Allocation a)  -- allocation map
+regalloc vars' regs aliaspairs =
+    let foldfunc = \st' ((int, name), index) ->
+            let st = expireOldIntervals st' int
+                wanted = findWantedAliases aliaspairs name
+                wantedregs = uniq $ sort $ catMaybes $ flip map wanted $ \n ->
+                                case findIndex (== n) intnames of
+                                    Nothing -> Nothing
+                                    Just idx | idx >= length (stAlloc st) -> Nothing
+                                             | otherwise -> case stAlloc st !! idx of
+                                        AllocMem -> Nothing
+                                        AllocReg r -> Just r
+            in if length (stActive st) == length regs
+                   then spillAtInterval st index
+                   else let -- ([regchoice], fr) = splitAt 1 (stFreeRegs st)
+                            (regchoice, fr) = case find (`elem` wantedregs) (stFreeRegs st) of
+                                Nothing -> (head (stFreeRegs st), tail (stFreeRegs st))
+                                Just wr -> trace ("Pair-allocated " ++ show name ++ " in " ++ show wr) $
+                                           (wr, stFreeRegs st \\ [wr])
+                            allocrev = stAlloc st ++ [AllocReg regchoice]
+                            active = sortBy (compare `on` snd . (ints !!)) $ index : stActive st
+                        in State active allocrev fr
+    in Map.fromList $ zip intnames $ stAlloc $
+               foldl foldfunc (State [] [] regs) (zip vars [0..])
+  where
+    vars = sortBy (compare `on` fst . fst) vars'
+    (ints, intnames) = (map fst vars, map snd vars)
+
+    expireOldIntervals :: State a -> Interval -> State a
+    expireOldIntervals st (intstart, _) =
+        let (dropped, active) = span ((< intstart) . snd . (ints !!)) (stActive st)
+            fr = selectAllocRegs (map (stAlloc st !!) dropped) ++ stFreeRegs st
+        in State active (stAlloc st) fr
+
+    spillAtInterval :: State a -> Int -> State a
+    spillAtInterval st index =
+        let spill = last (stActive st)
+        in if snd (ints !! spill) > snd (ints !! index)
+               then let alloc = setAt spill AllocMem (stAlloc st) ++ [stAlloc st !! spill]
+                        active = sortBy (compare `on` snd . (ints !!)) $
+                                     index : (stActive st \\ [spill])
+                    in State active alloc (stFreeRegs st)
+               else State (stActive st) (stAlloc st ++ [AllocMem]) (stFreeRegs st)
+
+    findWantedAliases :: (Ord b) => [(b, b)] -> b -> [b]
+    findWantedAliases pairs x =
+        uniq $ sort $ map snd (filter ((== x) . fst) pairs) ++ map fst (filter ((== x) . snd) pairs)
+
+selectAllocRegs :: [Allocation a] -> [a]
+selectAllocRegs allocs = catMaybes $ flip map allocs $ \alloc -> case alloc of
+        (AllocReg r) -> Just r
+        AllocMem -> Nothing
+
+setAt :: Int -> a -> [a] -> [a]
+setAt idx v l =
+    let (pre, _ : post) = splitAt idx l
+    in pre ++ v : post