module Compiler where

import Constant
import Control.Monad
import Data.Char (isUpper)
import qualified Data.Map as M

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

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 i ps) = CallI i $ map go ps
    go o@(ConstI (Atom i))
      | i == void = VoidI
      | varname (itos M.! i) = VarI i
      | otherwise = o
    go x = x

renumVars :: (Int -> Maybe PrlgInt) -> PrlgInt -> PrlgInt
renumVars rename = go
  where
    go (CallI i ps) = CallI i $ map go ps
    go (VarI idx)
      | 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') | ((idx, n), idx') <- zip occurs [1 ..], n > 1]
   in renumVars (m' M.!?) x

squashChoices :: [Code] -> Either String Code
squashChoices = out . concatMap go
  where
    go [Choices cs] = cs
    go x = [x]
    out [] = Left "goal compilation has no choices?"
    out [x] = pure x
    out xs = pure [Choices xs]

compileGoals :: Int -> Int -> Int -> PrlgInt -> Either String Code
compileGoals andop orop cut = fmap (++ [Done]) . go'
  where
    go' = struct2goal >=> go
    go p@(CallI x args@[_, _])
      | x == andop = concat <$> traverse go' args
      | x == orop = traverse go' args >>= squashChoices
    go p@(CallI x [])
      | x == cut = pure [Cut]
    go x = compileGoal x

compileGoal :: PrlgInt -> Either String Code
compileGoal = fmap compileArg . struct2goal

compileArg :: PrlgInt -> Code
compileArg (CallI i args) =
  U (Struct Id {str = i, arity = length args}) : concatMap compileArg args
compileArg (ConstI c) = [U (C c)]
compileArg (VarI x) = [U (LocalRef x)]
compileArg (VoidI) = [U VoidRef]

seqGoals :: [Code] -> Code
seqGoals = (++ [Done]) . concat

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

struct2goal :: PrlgInt -> Either String PrlgInt
struct2goal (ConstI (Atom s)) = pure $ CallI s []
struct2goal call@(CallI _ _) = pure call
struct2goal x = Left $ "cannot compile goal: " ++ show x