content: {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE FlexibleContexts #-}
  
  -- |
  
  module JL.Inferer where
  
  
  import           Control.Monad.State.Strict
  import qualified Data.HashMap.Strict as HM
  import           Data.Map (Map)
  import qualified Data.Map as M
  import           Data.Set (Set)
  import qualified Data.Set as S
  import           Data.Text (Text)
  import qualified Data.Text as T
  import           JL.Printer
  import           JL.Types
  
  -- | Get the type of the expression.
  infer :: Map Variable Type -> Expression -> [TypeVariable] -> Type
  infer ctx t stream =
    case evalState (check ctx t) stream of
      (ty, cs) ->
        let s = evalState (unify (S.toList cs)) ()
        in replace s ty
  
  -- | Check the exp with the given context.
  check
    :: MonadState ([TypeVariable]) m
    => Map Variable Type -> Expression -> m (Type, Set (Type, Type))
  check ctx expr =
    case expr of
      VariableExpression name@(Variable text) ->
        case M.lookup name ctx of
          Nothing -> error ("Not in scope: `" <> T.unpack text <> "'")
          Just typ -> return (typ, mempty)
      LambdaExpression x body -> do
        sym <- generateTypeVariable
        let xty = VariableType sym
        (rty, cs) <- check (M.insert x xty ctx) body
        return (FunctionType xty rty, cs)
      ConstantExpression {} -> return (JSONType, mempty)
      ApplicationExpression f x -> do
        (fty, cs1) <- check ctx f
        (xty, cs2) <- check ctx x
        sym <- generateTypeVariable
        let rty = VariableType sym
            cs = S.insert (fty, FunctionType xty rty) (cs1 <> cs2)
        return (rty, cs)
      InfixExpression l f r -> do
        (fty, cs1) <- check ctx (VariableExpression f)
        (ty1, cs2) <- check ctx l
        (ty2, cs3) <- check ctx r
        sym <- generateTypeVariable
        let rty = VariableType sym
            cs =
              S.insert
                (fty, FunctionType ty1 (FunctionType ty2 rty))
                (cs1 <> cs2 <> cs3)
        return (rty, cs)
      IfExpression cond a b -> do
        (condty, cs1) <- check ctx cond
        (aty, cs2) <- check ctx a
        (bty, cs3) <- check ctx b
        sym <- generateTypeVariable
        let rty = VariableType sym
            cs =
              S.insert
                (condty, JSONType)
                (S.insert (aty, bty) (cs1 <> cs2 <> cs3))
        pure (rty, cs)
      RecordExpression pairs -> do
        cs <-
          foldM
            (\cs (_, e) -> do
               (pty, cs') <- check ctx e
               pure (S.insert (pty, JSONType) (cs <> cs')))
            mempty
            (HM.toList pairs)
        pure (JSONType, cs)
      SubscriptExpression e ks -> do
        (t1, c1) <-
          (case e of
             ExpressionSubscripted es -> check ctx es
             WildcardSubscripted -> pure (FunctionType JSONType JSONType, mempty))
        cs <-
          foldM
            (\cs s ->
               case s of
                 PropertySubscript {} -> pure cs
                 ExpressionSubscript es -> do
                   (pty, cs') <- check ctx es
                   pure (S.insert (pty, JSONType) (cs <> cs')))
            c1
            ks
        let rty = case e of
                    WildcardSubscripted -> FunctionType JSONType JSONType
                    _ -> JSONType
        pure
          ( rty
          , S.insert (t1, rty) cs)
      ArrayExpression as -> do
        cs <-
          foldM
            (\cs e -> do
               (pty, cs') <- check ctx e
               pure (S.insert (pty, JSONType) (cs <> cs')))
            mempty
            as
        pure (JSONType, cs)
  
  -- | Generate a fresh type variable.
  generateTypeVariable
    :: MonadState ([TypeVariable]) m
    => m TypeVariable
  generateTypeVariable =
    get >>= \case
      v:vs -> do
        put vs
        pure v
      _ ->
        error "Ran out of type variables"
  
  -- | Unify the list of constraints.
  unify
    :: Monad m
    => [(Type, Type)] -> m (Map TypeVariable Type)
  unify [] = return mempty
  unify ((a, b):cs)
    | a == b = unify cs
    | VariableType v <- a = unifyVariable v cs a b
    | VariableType v <- b = unifyVariable v cs b a
    | FunctionType a1 b1 <- a
    , FunctionType a2 b2 <- b = unify ([(a1, a2), (b1, b2)] <> cs)
    | otherwise =
      error
        (T.unpack
           ("Type " <> quote (prettyType a) <> " doesn't match " <>
            quote (prettyType b)))
  
  -- | Unify a type variable for two types.
  unifyVariable
    :: Monad m
    => TypeVariable -> [(Type, Type)] -> Type -> Type -> m (Map TypeVariable Type)
  unifyVariable v cs a b =
    if occurs v b
      then error
             (T.unpack ("Occurs check: " <> prettyType a <> " ~ " <> prettyType b))
      else let subbed = M.singleton v b
           in do rest <- unify (substitute subbed cs)
                 return (rest <> subbed)
  
  -- | Occurs check.
  occurs :: TypeVariable -> Type -> Bool
  occurs x (VariableType y)
    | x == y = True
    | otherwise = False
  occurs x (FunctionType a b) = occurs x a || occurs x b
  occurs _ JSONType = False
  
  -- | Substitute the unified type into the constraints.
  substitute :: Map TypeVariable Type -> [(Type, Type)] -> [(Type, Type)]
  substitute subs = map go
    where
      go (a, b) = (replace subs a, replace subs b)
  
  -- | Apply a substitution to a type.
  replace :: Map TypeVariable Type -> Type -> Type
  replace s' t' = M.foldrWithKey go t' s'
    where
      go s1 t (VariableType s2)
        | s1 == s2 = t
        | otherwise = VariableType s2
      go s t (FunctionType t2 t3) = FunctionType (go s t t2) (go s t t3)
      go _ _ JSONType = JSONType
  
  -- | Quote something for showing to the programmer.
  quote :: Text -> Text
  quote t = "‘" <> t <> "’"
  