module Optimiser(optimise) where

import Data.List

-- import Debug.Trace

import AST


type Optimisation = [Instruction] -> [Instruction]


optimisations :: [Optimisation]
optimisations =
    [collectSimilar, nullOps, collectAdds, propagateKnowns, firstSets, propagateSlides, uselessEnd, specialLoops, normaliseOrder, deduplicateSets]

composedOpts :: Optimisation
composedOpts = foldl1 (.) (reverse optimisations)
-- composedOpts = foldl1 (.) (map (traceShowId .) $ reverse optimisations)

optimise :: Program -> Program
optimise (Program inss) =
    Program $ repeated composedOpts inss


collectSimilar :: Optimisation
collectSimilar [] = []
collectSimilar (IAdd x off1 : IAdd y off2 : rest) | off1 == off2 =
    collectSimilar $ IAdd (x + y) off1 : rest
collectSimilar (ISlide off1 : ISlide off2 : rest) =
    collectSimilar $ ISlide (off1 + off2) : rest
collectSimilar (ILoop inss off : rest) = ILoop (collectSimilar inss) off : collectSimilar rest
collectSimilar (ins : rest) = ins : collectSimilar rest

nullOps :: Optimisation
nullOps [] = []
nullOps (IAdd 0 _ : rest) = nullOps rest
nullOps (ISlide 0 : rest) = nullOps rest
nullOps (ILoop inss off : rest) = ILoop (nullOps inss) off : nullOps rest
nullOps (ins : rest) = ins : nullOps rest

collectAdds :: Optimisation
collectAdds [] = []
collectAdds inss =
    let adds = map (\(IAdd v o) -> (v, o)) $ takeWhile isIAdd inss
        dests = nub $ map snd adds
        collected = [(sum $ map fst $ filter ((== d) . snd) adds, d) | d <- dests]
        rest = drop (length adds) inss
        rest' = case rest of
            [] -> []
            (i:is) -> i : collectAdds is
    in map (uncurry IAdd) collected ++ rest'

propagateKnowns :: Optimisation
propagateKnowns [] = []
propagateKnowns (ISet v off : rest) =
    let pre = takeWhile (\ins -> isIAdd ins || isISet ins || isIMove ins) rest
        post = drop (length pre) rest
        relevant = filter ((== off) . offsetOf) pre
        irrelevant = filter ((/= off) . offsetOf) pre
        (res, resins) = accumSetAddMove v relevant
    in ISet res off : resins ++ propagateKnowns (irrelevant ++ post)
propagateKnowns (ILoop inss off : rest) = ILoop (propagateKnowns inss) off : propagateKnowns rest
propagateKnowns (ins : rest) = ins : propagateKnowns rest

firstSets :: Optimisation
firstSets [] = []
firstSets (IStart : rest) =
    let pre = takeWhile (\ins -> isIAdd ins || isISet ins) rest
        post = drop (length pre) rest
        dests = nub $ map offsetOf pre
        collected = [ISet (accumSetAdd 0 (filter ((== d) . offsetOf) pre)) d | d <- dests]
    in IStart : collected ++ post
firstSets inss = inss

propagateSlides :: Optimisation
propagateSlides [] = []
propagateSlides (ISlide off : rest) = propagateSlides (incOffsets off rest) ++ [ISlide off]
propagateSlides (ILoop inss off : rest) = ILoop (propagateSlides inss) off : propagateSlides rest
propagateSlides (ins : rest) = ins : propagateSlides rest

uselessEnd :: Optimisation
uselessEnd [] = []
uselessEnd inss = reverse $ dropWhile isUseless $ reverse inss
  where
    isUseless :: Instruction -> Bool
    isUseless (IInput _) = False
    isUseless (IOutput _) = False
    isUseless (ILoop lp _) = all isUseless lp
    isUseless _ = True

specialLoops :: Optimisation
specialLoops [] = []
specialLoops (ILoop [IAdd v off2] off1 : rest)
    | off1 /= off2 = ILoop [] off1 : specialLoops rest
    | gcd v 2 == 1 = ISet 0 off1 : specialLoops rest
specialLoops (ILoop inss off : rest)
    | all isIAdd inss,
      sum (map (\(IAdd v _) -> v) $ filter ((== off) . offsetOf) inss) == -1 =
        let others = map (\(IAdd v o) -> (v, o)) $ filter ((/= off) . offsetOf) inss
            dests = nub $ map snd others
            tos = [(d, sum $ map fst $ filter ((== d) . snd) others) | d <- dests]
        in IMove off tos : specialLoops rest
specialLoops (ILoop inss off : rest) = ILoop (specialLoops inss) off : specialLoops rest
specialLoops (ins : rest) = ins : specialLoops rest

normaliseOrder :: Optimisation
normaliseOrder [] = []
normaliseOrder inss =
    let pre = takeWhile (\ins -> isIAdd ins || isISet ins) inss
        post = drop (length pre) inss
    in if null pre
           then head inss : normaliseOrder (tail inss)
           else filter isISet pre ++ filter isIAdd pre ++ normaliseOrder post

deduplicateSets :: Optimisation
deduplicateSets [] = []
deduplicateSets (IStart : ISet 0 _ : rest) = IStart : deduplicateSets rest
deduplicateSets (ISet _ o1 : ISet v o2 : rest) | o1 == o2 = deduplicateSets $ ISet v o1 : rest
deduplicateSets (ins : rest) = ins : deduplicateSets rest


repeated :: Optimisation -> Optimisation
repeated opt = \inss -> let inss' = opt inss
                        in if inss == inss' then inss else repeated opt inss'

incOffsets :: Offset -> [Instruction] -> [Instruction]
incOffsets _ [] = []
incOffsets inc (IAdd v off : rest) = IAdd v (off + inc) : incOffsets inc rest
incOffsets inc (ISet v off : rest) = ISet v (off + inc) : incOffsets inc rest
incOffsets inc (IMove from tos : rest) = IMove (from + inc) [(o+inc,m) | (o,m) <- tos] : incOffsets inc rest
incOffsets inc (ISlide off : rest) = ISlide off : incOffsets inc rest
incOffsets inc (ILoop inss off : rest) = ILoop (incOffsets inc inss) (off + inc) : incOffsets inc rest
incOffsets inc (IInput off : rest) = IInput (off + inc) : incOffsets inc rest
incOffsets inc (IOutput off : rest) = IOutput (off + inc) : incOffsets inc rest
incOffsets inc (IStart : rest) = IStart : incOffsets inc rest

accumSetAdd :: Byte -> [Instruction] -> Byte
accumSetAdd acc [] = acc
accumSetAdd _ (ISet v' _ : rest) = accumSetAdd v' rest
accumSetAdd acc (IAdd v' _ : rest) = accumSetAdd (acc + v') rest
accumSetAdd _ _ = unreachable

accumSetAddMove :: Byte -> [Instruction] -> (Byte, [Instruction])
accumSetAddMove acc [] = (acc, [])
accumSetAddMove _ (ISet v' _ : rest) = accumSetAddMove v' rest
accumSetAddMove acc (IAdd v' _ : rest) = accumSetAddMove (acc + v') rest
accumSetAddMove acc (IMove _ tos : rest) =
    let (res, resins) = accumSetAddMove 0 rest
    in (res, [IAdd (m * acc) o | (o,m) <- tos] ++ resins)
accumSetAddMove _ _ = unreachable

unreachable :: a
unreachable = error "Unreachable"