{-# LANGUAGE ScopedTypeVariables, TypeApplications, FlexibleContexts, AllowAmbiguousTypes #-}
module Overload.Normal where

import Control.Effects.State
import Data.Functor.Identity

import Overload.TypeTree

newtype Normal = Normal Int deriving (Eq, Ord)
instance Show Normal where
    show (Normal n) = show n

type FreshSource a = ([(a, Normal)], Normal)

lookupName :: (MonadEffectState (FreshSource a) m, Eq a) => a -> m Normal
lookupName name = do
    (table, Normal top) <- getState
    case lookup name table of
        Nothing -> do
            setState ((name, Normal top) : table, Normal (top + 1))
            return (Normal top)
        Just n  -> return n

freshVar :: forall a m. MonadEffectState (FreshSource a) m => m Normal
freshVar = do
    (table, Normal top) :: FreshSource a <- getState
    setState (table, Normal (top + 1))
    return (Normal top)

normalizeTypeTree :: forall a. Eq a => TypeTree (Maybe a) -> TypeTree Normal
normalizeTypeTree =
      runIdentity
    . handleStateT (([], Normal 0) :: FreshSource a)
    . traverse (maybe (freshVar @a) lookupName)