import Control.Monad.Fix

type Register = Int
type RegFile  = [Register]
type Label    = String

data Instruction = ZERO  Register 
                 | ADD   Register Int
                 | MOV   Register Register
                 | LABEL Label
                 | JIP   Register Label
                 deriving Show

type Program = [Instruction]

class MonadFix m => Processor m where
   zero       :: Register -> m ()
   mov        :: Register -> Register -> m ()
   addi       :: Register -> Int -> m ()
   label      :: m Label
   brPositive :: Register -> Label -> m ()
-----------------------------------------------------------------------
data Interpreter a = Interp ((Int, Program) -> (a, (Int, Program)))

instance Monad Interpreter where
   return x       = Interp (\s -> (x, s))
   Interp f >>= g = Interp (\s -> let (a, s')  = f s
                                      Interp h = g a
                                  in h s')

instance MonadFix Interpreter where
   mfix f = Interp (\s -> let Interp h = f a 
                              (a, s')  = h s
                          in (a, s'))

mkLabel :: Int -> Label
mkLabel i = 'L' : show i 

newLabel :: Interpreter Label
newLabel = Interp (\(i, p) -> (mkLabel i, (i+1, p)))

-- note: we build programs backwards! (for efficiency)
emit :: Instruction -> Interpreter ()
emit ins = Interp (\(i, p) -> ((), (i, ins : p)))

instance Processor Interpreter where
   zero r    = emit (ZERO r)
   addi r i  = emit (ADD r i)
   mov r1 r2 = emit (MOV r1 r2)
   label     = do l <- newLabel
                  emit (LABEL l)
                  return l
   brPositive r l = emit (JIP r l)
------------------------------------------------------------------------
updateRF :: RegFile -> Register -> (Int -> Int) -> RegFile
updateRF rf r f = p ++ f v : s
          where (p, v : s) = splitAt r (rf ++ pad)
                pad = take (r - length rf + 1) (repeat 0)

lookUpRegFile :: RegFile -> Register -> Int
lookUpRegFile rf r = (rf ++ repeat 0) !! r

execute :: Program -> RegFile
execute pgm = step 0 []
  where endOfProg = length pgm
        step ic rf | ic == endOfProg = rf
        step ic rf = let inst = pgm !! ic
                     in step (next inst rf ic) (doInst inst rf)
        next (JIP r i) rf ic 
          | lookUpRegFile rf r > 0 = find i pgm 0
          | True                   = ic + 1
        next _         _  ic       = ic + 1
        doInst (ZERO r)    rf = updateRF rf r (const 0)
        doInst (ADD r i)   rf = updateRF rf r (+i)
        doInst (MOV r1 r2) rf = updateRF rf r1 (const (lookUpRegFile rf r2))
        doInst _           rf = rf
        find i []           _ = error ("jump to unknown label: " ++ show i)
        find i (LABEL j:is) k 
          | i == j = k
          | True   = find i is (k+1)
        find i (_:is) k = find i is (k+1)
------------------------------------------------------------------------
compile :: Interpreter () -> Program
compile (Interp f) = let (_, (_, pgm)) = f (0, []) in reverse pgm

run :: Interpreter () -> RegFile
run = execute . compile 
------------------------------------------------------------------------
-- examples:

r0, r1, r2, r3 :: Register
(r0 : r1 : r2 : r3 : _) = [0 .. ]

compute112 :: Processor m => m ()
compute112 = do
        zero r1
        zero r2
        addi r2 3
        topLoop <- label
        addi r1 4
        addi r2 (-1)
        brPositive r2 topLoop
        addi r1 100

compute12 :: Processor m => m ()
compute12 = mdo 
        zero r1 
        zero r2
        addi r2 3        -- loop three times
        topLoop <- label
        addi r1 4
        addi r2 (-1)        -- decrement loop counter
        brPositive r2 topLoop
        mov r3 r1
        addi r3 (-10)        
        brPositive r3 out -- if r1 > 10 don't add 100
        addi r1 100
        out <- label
        return ()
------------------------------------------------------------------------
test1 = compile compute112
test2 = run compute112     -- yields [0, 112, 0], i.e. r0 = 0, r1 = 112
test3 = compile compute12 
test4 = run compute12      -- yields [0, 12, 0, 2], i.e. r0 = 0, r1 = 12