module Compiler where

import Data.Char (isUpper)
import qualified Data.Map as M

import Code (Code, Datum(..), Heap, Instr(..), heapStruct)
import IR (Id(..), PrlgInt(..), StrTable(..))

desugarPrlg :: Int -> PrlgInt -> PrlgInt
desugarPrlg list = go
  where
    go (CallI id ps) = CallI id $ map go ps
    go (ListI (x:xs) t) =
      CallI Id {str = list, arity = 2} [go x, go (ListI xs t)]
    go (ListI [] Nothing) = LiteralI list
    go (ListI [] (Just x)) = go x
    go x = x

varname :: String -> Bool
varname ('_':_) = True
varname (c:_) = isUpper c
varname _ = False

varOccurs :: PrlgInt -> M.Map Int Int
varOccurs (CallI _ xs) = M.unionsWith (+) $ map varOccurs xs
varOccurs (VarI idx _) = M.singleton idx 1
varOccurs _ = M.empty

variablizePrlg :: Int -> StrTable -> PrlgInt -> PrlgInt
variablizePrlg void (StrTable _ _ itos) = go
  where
    go (CallI id ps) = CallI id $ map go ps
    go (LiteralI i)
      | i == void = VoidI
      | varname (itos M.! i) = VarI i i
      | otherwise = LiteralI i

renumVars :: (Int -> Maybe PrlgInt) -> PrlgInt -> PrlgInt
renumVars rename = go
  where
    go (CallI id ps) = CallI id $ map go ps
    go (VarI idx i)
      | Just new <- rename idx = new
    go x = x

squashVars :: PrlgInt -> PrlgInt
squashVars x =
  let occurs = M.assocs (varOccurs x)
      m' =
        M.fromList $
        [(idx, VoidI) | (idx, n) <- occurs, n <= 1] ++
        [(idx, VarI idx' 0) | ((idx, n), idx') <- zip occurs [1 ..], n > 1]
   in renumVars (m' M.!?) x

compileGoals :: Id -> Int -> PrlgInt -> [Code]
compileGoals andop cut = go'
  where
    go' = go . struct2goal
    go p@(CallI x args)
      | x == andop = concatMap go' args
    go p@(CallI (Id x 0) [])
      | x == cut = [[Cut]]
    go x = [compileGoal x]

compileGoal :: PrlgInt -> Code
compileGoal = compileArg . struct2goal

compileArg :: PrlgInt -> Code
compileArg (CallI s args) = U (Struct s) : concatMap compileArg args
compileArg (LiteralI s) = [U (Atom s)]
compileArg (VarI x s) = [U (LocalRef x s)]
compileArg (VoidI) = [U VoidRef]

seqGoals :: [Code] -> Code
seqGoals [] = [NoGoal]
seqGoals [[Cut]] = [Cut, NoGoal]
seqGoals [x] = [Goal] ++ x ++ [LastCall]
seqGoals [x, [Cut]] = [Goal] ++ x ++ [LastCall, Cut]
seqGoals ([Cut]:xs) = [Cut] ++ seqGoals xs
seqGoals (x:xs) = [Goal] ++ x ++ [Call] ++ seqGoals xs

heapStructPrlgInt :: Monad m => m PrlgInt -> Heap -> Int -> m PrlgInt
heapStructPrlgInt heaperr heap ref = heapStruct atom struct hrec heap ref
  where
    atom (Atom s) = pure $ LiteralI s
    atom VoidRef = pure $ VoidI
    struct (Struct s) args = pure $ CallI s args
    hrec (HeapRef r) ref
      | r == ref = pure $ VarI r 0
      | otherwise = heaperr

-- TODO check if this is used
goal2struct :: PrlgInt -> PrlgInt
goal2struct (CallI (Id s 0) []) = LiteralI s
goal2struct x = x

struct2goal :: PrlgInt -> PrlgInt
struct2goal (LiteralI s) = CallI (Id s 0) []
struct2goal x = x