elm/compiler/Type/Environment.hs

module Type.Environment where

import Control.Applicative ((<$>), (<*>))
import Control.Monad
import qualified Control.Monad.State as State
import qualified Data.Traversable as Traverse
import qualified Data.Map as Map
import Data.List (isPrefixOf)
import qualified Data.UnionFind.IO as UF

import qualified SourceSyntax.Type as Src
import SourceSyntax.Module (ADT)
import Type.Type

type TypeDict = Map.Map String Type
type VarDict = Map.Map String Variable

data Environment = Environment {
  constructor :: Map.Map String (IO (Int, [Variable], [Type], Type)),
  aliases :: Map.Map String ([String], Src.Type),
  types :: TypeDict,
  value :: TypeDict
}

initialEnvironment :: [ADT] -> [(String, [String], Src.Type)] -> IO Environment
initialEnvironment datatypes aliases = do
    types <- makeTypes datatypes
    let aliases' = Map.fromList $ map (\(a,b,c) -> (a,(b,c))) aliases
        env = Environment {
                constructor = Map.empty,
                value = Map.empty,
                types = types,
                aliases = aliases' }

    return $ env { constructor = makeConstructors env datatypes }

makeTypes :: [ADT] -> IO TypeDict
makeTypes datatypes = 
    Map.fromList <$> mapM makeCtor (builtins ++ map nameAndKind datatypes)
  where
    nameAndKind (name, tvars, _) = (name, length tvars)

    makeCtor (name, kind) = do
      ctor <- VarN <$> namedVar Constant name
      return (name, ctor)

    tuple n = ("_Tuple" ++ show n, n)

    kind n names = map (\name -> (name, n)) names

    builtins :: [(String,Int)]
    builtins = concat [ map tuple [0..9]
                      , kind 1 ["_List"]
                      , kind 0 ["Int","Float","Char","Bool"]
                      ]


makeConstructors :: Environment
                 -> [ADT]
                 -> Map.Map String (IO (Int, [Variable], [Type], Type))
makeConstructors env datatypes = Map.fromList builtins
  where
    list t = (types env Map.! "_List") <| t

    inst :: Int -> ([Type] -> ([Type], Type)) -> IO (Int, [Variable], [Type], Type)
    inst numTVars tipe = do
      vars <- forM [1..numTVars] $ \_ -> var Flexible
      let (args, result) = tipe (map VarN vars)
      return (length args, vars, args, result)

    tupleCtor n =
        let name = "_Tuple" ++ show n
        in  (name, inst n $ \vs -> (vs, foldl (<|) (types env Map.! name) vs))
    
    builtins :: [ (String, IO (Int, [Variable], [Type], Type)) ]
    builtins = [ ("[]", inst 1 $ \ [t] -> ([], list t))
               , ("::", inst 1 $ \ [t] -> ([t, list t], list t))
               ] ++ map tupleCtor [0..9]
                 ++ concatMap (ctorToType env) datatypes


ctorToType :: Environment -> ADT -> [ (String, IO (Int, [Variable], [Type], Type)) ]
ctorToType env (name, tvars, ctors) =
    zip (map fst ctors) (map inst ctors)
  where
    inst :: (String, [Src.Type]) -> IO (Int, [Variable], [Type], Type)
    inst ctor = do
      ((args, tipe), (dict,_)) <- State.runStateT (go ctor) (Map.empty, Map.empty)
      return (length args, Map.elems dict, args, tipe)
      

    go :: (String, [Src.Type]) -> State.StateT (VarDict, TypeDict) IO ([Type], Type)
    go (ctor, args) = do
      types <- mapM (instantiator env) args
      returnType <- instantiator env (Src.Data name (map Src.Var tvars))
      return (types, returnType)


get :: Environment -> (Environment -> Map.Map String a) -> String -> a
get env subDict key = Map.findWithDefault err key (subDict env)
  where
    err = error $ "\nCould not find type constructor '" ++ key ++ "' while checking types."


freshDataScheme :: Environment -> String -> IO (Int, [Variable], [Type], Type)
freshDataScheme env name = get env constructor name

instantiateType ::
    Environment -> Src.Type -> VarDict -> IO ([Variable], Type)
instantiateType env sourceType dict =
  do (tipe, (dict',_)) <- State.runStateT (instantiator env sourceType) (dict, Map.empty)
     return (Map.elems dict', tipe)

instantiator :: Environment -> Src.Type
             -> State.StateT (VarDict, TypeDict) IO Type
instantiator env sourceType = go sourceType
  where
    go :: Src.Type -> State.StateT (VarDict, TypeDict) IO Type
    go sourceType =
      case sourceType of
        Src.Lambda t1 t2 -> (==>) <$> go t1 <*> go t2

        Src.Var x -> do
          (dict, aliases) <- State.get
          case Map.lookup x dict of
            Just var -> return (VarN var)
            Nothing ->
                case Map.lookup x aliases of
                  Just t -> return t
                  Nothing ->
                    do var <- State.liftIO $ namedVar flex x
                       State.put (Map.insert x var dict, aliases)
                       return (VarN var)
                    where
                      flex | "number"     `isPrefixOf` x = Is Number
                           | "comparable" `isPrefixOf` x = Is Comparable
                           | "appendable" `isPrefixOf` x = Is Appendable
                           | otherwise = Flexible

        Src.Data "String" [] ->
            return (get env types "_List" <| get env types "Char")

        Src.Data name ts -> do
          ts' <- mapM go ts
          case Map.lookup name (types env) of
            Just t -> return $ foldl (<|) t ts'
            Nothing ->
                case Map.lookup name (aliases env) of
                  Nothing -> error $ "\nCould not find type constructor '" ++ name ++ "' while checking types."
                  Just (tvars, t) ->
                      let tvarLen = length tvars
                          msg = "\nType alias '" ++ name ++ "' expects " ++ show tvarLen ++
                                " type argument" ++ (if tvarLen == 1 then "" else "s") ++
                                " but was given " ++ show (length ts')
                      in  if length ts' /= length tvars then error msg else
                              do (dict, aliases) <- State.get
                                 State.put (dict, Map.union aliases . Map.fromList $ zip tvars ts')
                                 go t

        Src.EmptyRecord -> return (TermN EmptyRecord1)

        Src.Record fields ext ->
          TermN <$> (Record1 <$> Traverse.traverse (mapM go) (Src.fieldMap fields) <*> go ext)