module IR where

import Data.Char (isNumber)
import Data.List (mapAccumL)
import qualified Data.Map as M

data PrlgStr
  = CallS String [PrlgStr]
  | LiteralS String
  | ListS [PrlgStr] (Maybe PrlgStr)
  deriving (Show)

data Id =
  Id
    { str :: Int
    , arity :: Int
    }
  deriving (Show, Eq, Ord)

data PrlgInt
  = CallI Int [PrlgInt]
  | AtomI Int
  | NumI Int
  | ListI [PrlgInt] (Maybe PrlgInt) -- only exists before desugaring
  | VarI Int Int -- VarI localIndex strTableString
  | VoidI
  deriving (Show)

data StrTable =
  StrTable Int (M.Map String Int) (M.Map Int String)
  deriving (Show)

emptystrtable = StrTable 1 M.empty M.empty

strtablize t@(StrTable nxt fwd rev) str =
  case fwd M.!? str of
    Just i -> (t, i)
    _ -> (StrTable (nxt + 1) (M.insert str nxt fwd) (M.insert nxt str rev), nxt)

internPrlg :: StrTable -> PrlgStr -> (StrTable, PrlgInt)
internPrlg = go
  where
    go t (LiteralS str)
      | all isNumber str = (t, NumI $ read str)
      | otherwise = AtomI <$> strtablize t str
    go t (CallS str ps) =
      let (t', i) = strtablize t str
       in CallI i <$> mapAccumL go t' ps
    go t (ListS elems Nothing) = flip ListI Nothing <$> mapAccumL go t elems
    go t (ListS elems (Just tail)) =
      let (t', tail') = go t tail
       in flip ListI (Just tail') <$> mapAccumL go t' elems