content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE Rank2Types #-}
  {-# OPTIONS_GHC -Wall #-}
  
  module Type.Unify
    ( Answer (..),
      unify,
    )
  where
  
  import Data.Map.Strict qualified as Map
  import Data.Name qualified as Name
  import Gren.ModuleName qualified as ModuleName
  import Type.Error qualified as Error
  import Type.Occurs qualified as Occurs
  import Type.Type as Type
  import Type.UnionFind qualified as UF
  
  -- UNIFY
  
  data Answer
    = Ok [Variable]
    | Err [Variable] Error.Type Error.Type
  
  unify :: Variable -> Variable -> IO Answer
  unify v1 v2 =
    case guardedUnify v1 v2 of
      Unify k ->
        k [] onSuccess $ \vars () ->
          do
            t1 <- Type.toErrorType v1
            t2 <- Type.toErrorType v2
            UF.union v1 v2 errorDescriptor
            return (Err vars t1 t2)
  
  onSuccess :: [Variable] -> () -> IO Answer
  onSuccess vars () =
    return (Ok vars)
  
  errorDescriptor :: Descriptor
  errorDescriptor =
    Descriptor Error noRank noMark Nothing
  
  -- CPS UNIFIER
  
  newtype Unify a
    = Unify
        ( forall r.
          [Variable] ->
          ([Variable] -> a -> IO r) ->
          ([Variable] -> () -> IO r) ->
          IO r
        )
  
  instance Functor Unify where
    fmap func (Unify kv) =
      Unify $ \vars ok err ->
        let ok1 vars1 value =
              ok vars1 (func value)
         in kv vars ok1 err
  
  instance Applicative Unify where
    pure a =
      Unify $ \vars ok _ ->
        ok vars a
  
    (<*>) (Unify kf) (Unify kv) =
      Unify $ \vars ok err ->
        let ok1 vars1 func =
              let ok2 vars2 value =
                    ok vars2 (func value)
               in kv vars1 ok2 err
         in kf vars ok1 err
  
    (*>) (Unify ka) (Unify kb) =
      Unify $ \vars ok err ->
        let ok1 vars1 _ = kb vars1 ok err
         in ka vars ok1 err
  
  instance Monad Unify where
    (>>=) (Unify ka) callback =
      Unify $ \vars ok err ->
        let ok1 vars1 a =
              case callback a of
                Unify kb -> kb vars1 ok err
         in ka vars ok1 err
  
  register :: IO Variable -> Unify Variable
  register mkVar =
    Unify $ \vars ok _ ->
      do
        var <- mkVar
        ok (var : vars) var
  
  mismatch :: Unify a
  mismatch =
    Unify $ \vars _ err ->
      err vars ()
  
  -- UNIFICATION HELPERS
  
  data Context = Context
    { _first :: Variable,
      _firstDesc :: Descriptor,
      _second :: Variable,
      _secondDesc :: Descriptor
    }
  
  reorient :: Context -> Context
  reorient (Context var1 desc1 var2 desc2) =
    Context var2 desc2 var1 desc1
  
  -- MERGE
  
  merge :: Context -> Content -> Unify ()
  merge (Context var1 (Descriptor _ rank1 _ _) var2 (Descriptor _ rank2 _ _)) content =
    Unify $ \vars ok _ ->
      ok vars
        =<< UF.union var1 var2 (Descriptor content (min rank1 rank2) noMark Nothing)
  
  fresh :: Context -> Content -> Unify Variable
  fresh (Context _ (Descriptor _ rank1 _ _) _ (Descriptor _ rank2 _ _)) content =
    register $
      UF.fresh $
        Descriptor content (min rank1 rank2) noMark Nothing
  
  -- ACTUALLY UNIFY THINGS
  
  guardedUnify :: Variable -> Variable -> Unify ()
  guardedUnify left right =
    Unify $ \vars ok err ->
      do
        equivalent <- UF.equivalent left right
        if equivalent
          then ok vars ()
          else do
            leftDesc <- UF.get left
            rightDesc <- UF.get right
            case actuallyUnify (Context left leftDesc right rightDesc) of
              Unify k ->
                k vars ok err
  
  subUnify :: Variable -> Variable -> Unify ()
  subUnify var1 var2 =
    guardedUnify var1 var2
  
  actuallyUnify :: Context -> Unify ()
  actuallyUnify context@(Context _ (Descriptor firstContent _ _ _) _ (Descriptor secondContent _ _ _)) =
    case firstContent of
      FlexVar _ ->
        unifyFlex context firstContent secondContent
      FlexSuper super _ ->
        unifyFlexSuper context super firstContent secondContent
      RigidVar _ ->
        unifyRigid context Nothing firstContent secondContent
      RigidSuper super _ ->
        unifyRigid context (Just super) firstContent secondContent
      Alias home name args realVar ->
        unifyAlias context home name args realVar secondContent
      Structure flatType ->
        unifyStructure context flatType firstContent secondContent
      Error ->
        -- If there was an error, just pretend it is okay. This lets us avoid
        -- "cascading" errors where one problem manifests as multiple message.
        merge context Error
  
  -- UNIFY FLEXIBLE VARIABLES
  
  unifyFlex :: Context -> Content -> Content -> Unify ()
  unifyFlex context content otherContent =
    case otherContent of
      Error ->
        merge context Error
      FlexVar maybeName ->
        merge context $
          case maybeName of
            Nothing ->
              content
            Just _ ->
              otherContent
      FlexSuper _ _ ->
        merge context otherContent
      RigidVar _ ->
        merge context otherContent
      RigidSuper _ _ ->
        merge context otherContent
      Alias _ _ _ _ ->
        merge context otherContent
      Structure _ ->
        merge context otherContent
  
  -- UNIFY RIGID VARIABLES
  
  unifyRigid :: Context -> Maybe SuperType -> Content -> Content -> Unify ()
  unifyRigid context maybeSuper content otherContent =
    case otherContent of
      FlexVar _ ->
        merge context content
      FlexSuper otherSuper _ ->
        case maybeSuper of
          Just super ->
            if combineRigidSupers super otherSuper
              then merge context content
              else mismatch
          Nothing ->
            mismatch
      RigidVar _ ->
        mismatch
      RigidSuper _ _ ->
        mismatch
      Alias _ _ _ _ ->
        mismatch
      Structure _ ->
        mismatch
      Error ->
        merge context Error
  
  -- UNIFY SUPER VARIABLES
  
  unifyFlexSuper :: Context -> SuperType -> Content -> Content -> Unify ()
  unifyFlexSuper context super content otherContent =
    case otherContent of
      Structure flatType ->
        unifyFlexSuperStructure context super flatType
      RigidVar _ ->
        mismatch
      RigidSuper otherSuper _ ->
        if combineRigidSupers otherSuper super
          then merge context otherContent
          else mismatch
      FlexVar _ ->
        merge context content
      FlexSuper otherSuper _ ->
        case super of
          Number ->
            case otherSuper of
              Number -> merge context content
              Comparable -> merge context content
              Appendable -> mismatch
              CompAppend -> mismatch
          Comparable ->
            case otherSuper of
              Comparable -> merge context otherContent
              Number -> merge context otherContent
              Appendable -> merge context (Type.unnamedFlexSuper CompAppend)
              CompAppend -> merge context otherContent
          Appendable ->
            case otherSuper of
              Appendable -> merge context otherContent
              Comparable -> merge context (Type.unnamedFlexSuper CompAppend)
              CompAppend -> merge context otherContent
              Number -> mismatch
          CompAppend ->
            case otherSuper of
              Comparable -> merge context content
              Appendable -> merge context content
              CompAppend -> merge context content
              Number -> mismatch
      Alias _ _ _ realVar ->
        subUnify (_first context) realVar
      Error ->
        merge context Error
  
  combineRigidSupers :: SuperType -> SuperType -> Bool
  combineRigidSupers rigid flex =
    rigid == flex
      || (rigid == Number && flex == Comparable)
      || (rigid == CompAppend && (flex == Comparable || flex == Appendable))
  
  atomMatchesSuper :: SuperType -> ModuleName.Canonical -> Name.Name -> Bool
  atomMatchesSuper super home name =
    case super of
      Number ->
        isNumber home name
      Comparable ->
        isNumber home name
          || Error.isString home name
          || Error.isChar home name
      Appendable ->
        Error.isString home name
      CompAppend ->
        Error.isString home name
  
  isNumber :: ModuleName.Canonical -> Name.Name -> Bool
  isNumber home name =
    home == ModuleName.basics
      && (name == Name.int || name == Name.float)
  
  unifyFlexSuperStructure :: Context -> SuperType -> FlatType -> Unify ()
  unifyFlexSuperStructure context super flatType =
    case flatType of
      App1 home name [] ->
        if atomMatchesSuper super home name
          then merge context (Structure flatType)
          else mismatch
      App1 home name [variable] | home == ModuleName.array && name == Name.array ->
        case super of
          Number ->
            mismatch
          Appendable ->
            merge context (Structure flatType)
          Comparable ->
            do
              comparableOccursCheck context
              unifyComparableRecursive variable
              merge context (Structure flatType)
          CompAppend ->
            do
              comparableOccursCheck context
              unifyComparableRecursive variable
              merge context (Structure flatType)
      _ ->
        mismatch
  
  -- TODO: is there some way to avoid doing this?
  -- Do type classes require occurs checks?
  comparableOccursCheck :: Context -> Unify ()
  comparableOccursCheck (Context _ _ var _) =
    Unify $ \vars ok err ->
      do
        hasOccurred <- Occurs.occurs var
        if hasOccurred
          then err vars ()
          else ok vars ()
  
  unifyComparableRecursive :: Variable -> Unify ()
  unifyComparableRecursive var =
    do
      compVar <- register $
        do
          (Descriptor _ rank _ _) <- UF.get var
          UF.fresh $ Descriptor (Type.unnamedFlexSuper Comparable) rank noMark Nothing
      guardedUnify compVar var
  
  -- UNIFY ALIASES
  
  unifyAlias :: Context -> ModuleName.Canonical -> Name.Name -> [(Name.Name, Variable)] -> Variable -> Content -> Unify ()
  unifyAlias context home name args realVar otherContent =
    case otherContent of
      FlexVar _ ->
        merge context (Alias home name args realVar)
      FlexSuper _ _ ->
        subUnify realVar (_second context)
      RigidVar _ ->
        subUnify realVar (_second context)
      RigidSuper _ _ ->
        subUnify realVar (_second context)
      Alias otherHome otherName otherArgs otherRealVar ->
        if name == otherName && home == otherHome
          then Unify $ \vars ok err ->
            let ok1 vars1 () =
                  case merge context otherContent of
                    Unify k ->
                      k vars1 ok err
             in unifyAliasArgs vars context args otherArgs ok1 err
          else subUnify realVar otherRealVar
      Structure _ ->
        subUnify realVar (_second context)
      Error ->
        merge context Error
  
  unifyAliasArgs :: [Variable] -> Context -> [(Name.Name, Variable)] -> [(Name.Name, Variable)] -> ([Variable] -> () -> IO r) -> ([Variable] -> () -> IO r) -> IO r
  unifyAliasArgs vars context args1 args2 ok err =
    case args1 of
      (_, arg1) : others1 ->
        case args2 of
          (_, arg2) : others2 ->
            case subUnify arg1 arg2 of
              Unify k ->
                k
                  vars
                  (\vs () -> unifyAliasArgs vs context others1 others2 ok err)
                  (\vs () -> unifyAliasArgs vs context others1 others2 err err)
          _ ->
            err vars ()
      [] ->
        case args2 of
          [] ->
            ok vars ()
          _ ->
            err vars ()
  
  -- UNIFY STRUCTURES
  
  unifyStructure :: Context -> FlatType -> Content -> Content -> Unify ()
  unifyStructure context flatType content otherContent =
    case otherContent of
      FlexVar _ ->
        merge context content
      FlexSuper super _ ->
        unifyFlexSuperStructure (reorient context) super flatType
      RigidVar _ ->
        mismatch
      RigidSuper _ _ ->
        mismatch
      Alias _ _ _ realVar ->
        subUnify (_first context) realVar
      Structure otherFlatType ->
        case (flatType, otherFlatType) of
          (App1 home name args, App1 otherHome otherName otherArgs) | home == otherHome && name == otherName ->
            Unify $ \vars ok err ->
              let ok1 vars1 () =
                    case merge context otherContent of
                      Unify k ->
                        k vars1 ok err
               in unifyArgs vars context args otherArgs ok1 err
          (Fun1 arg1 res1, Fun1 arg2 res2) ->
            do
              subUnify arg1 arg2
              subUnify res1 res2
              merge context otherContent
          (EmptyRecord1, EmptyRecord1) ->
            merge context otherContent
          (Record1 fields ext, EmptyRecord1)
            | Map.null fields ->
                subUnify ext (_second context)
          (EmptyRecord1, Record1 fields ext)
            | Map.null fields ->
                subUnify (_first context) ext
          (Record1 fields1 ext1, Record1 fields2 ext2) ->
            Unify $ \vars ok err ->
              do
                structure1 <- gatherFields fields1 ext1
                structure2 <- gatherFields fields2 ext2
                case unifyRecord context structure1 structure2 of
                  Unify k ->
                    k vars ok err
          _ ->
            mismatch
      Error ->
        merge context Error
  
  -- UNIFY ARGS
  
  unifyArgs :: [Variable] -> Context -> [Variable] -> [Variable] -> ([Variable] -> () -> IO r) -> ([Variable] -> () -> IO r) -> IO r
  unifyArgs vars context args1 args2 ok err =
    case args1 of
      arg1 : others1 ->
        case args2 of
          arg2 : others2 ->
            case subUnify arg1 arg2 of
              Unify k ->
                k
                  vars
                  (\vs () -> unifyArgs vs context others1 others2 ok err)
                  (\vs () -> unifyArgs vs context others1 others2 err err)
          _ ->
            err vars ()
      [] ->
        case args2 of
          [] ->
            ok vars ()
          _ ->
            err vars ()
  
  -- UNIFY RECORDS
  
  unifyRecord :: Context -> RecordStructure -> RecordStructure -> Unify ()
  unifyRecord context (RecordStructure fields1 ext1) (RecordStructure fields2 ext2) =
    let sharedFields = Map.intersectionWith (,) fields1 fields2
        uniqueFields1 = Map.difference fields1 fields2
        uniqueFields2 = Map.difference fields2 fields1
     in if Map.null uniqueFields1
          then
            if Map.null uniqueFields2
              then do
                subUnify ext1 ext2
                unifySharedFields context sharedFields Map.empty ext1
              else do
                subRecord <- fresh context (Structure (Record1 uniqueFields2 ext2))
                subUnify ext1 subRecord
                unifySharedFields context sharedFields Map.empty subRecord
          else
            if Map.null uniqueFields2
              then do
                subRecord <- fresh context (Structure (Record1 uniqueFields1 ext1))
                subUnify subRecord ext2
                unifySharedFields context sharedFields Map.empty subRecord
              else do
                let otherFields = Map.union uniqueFields1 uniqueFields2
                ext <- fresh context Type.unnamedFlexVar
                sub1 <- fresh context (Structure (Record1 uniqueFields1 ext))
                sub2 <- fresh context (Structure (Record1 uniqueFields2 ext))
                subUnify ext1 sub2
                subUnify sub1 ext2
                unifySharedFields context sharedFields otherFields ext
  
  unifySharedFields :: Context -> Map.Map Name.Name (Variable, Variable) -> Map.Map Name.Name Variable -> Variable -> Unify ()
  unifySharedFields context sharedFields otherFields ext =
    do
      matchingFields <- Map.traverseMaybeWithKey unifyField sharedFields
      if Map.size sharedFields == Map.size matchingFields
        then merge context (Structure (Record1 (Map.union matchingFields otherFields) ext))
        else mismatch
  
  unifyField :: Name.Name -> (Variable, Variable) -> Unify (Maybe Variable)
  unifyField _ (actual, expected) =
    Unify $ \vars ok _ ->
      case subUnify actual expected of
        Unify k ->
          k
            vars
            (\vs () -> ok vs (Just actual))
            (\vs () -> ok vs Nothing)
  
  -- GATHER RECORD STRUCTURE
  
  data RecordStructure = RecordStructure
    { _fields :: Map.Map Name.Name Variable,
      _extension :: Variable
    }
  
  gatherFields :: Map.Map Name.Name Variable -> Variable -> IO RecordStructure
  gatherFields fields variable =
    do
      (Descriptor content _ _ _) <- UF.get variable
      case content of
        Structure (Record1 subFields subExt) ->
          gatherFields (Map.union fields subFields) subExt
        Alias _ _ _ var ->
          -- TODO may be dropping useful alias info here
          gatherFields fields var
        _ ->
          return (RecordStructure fields variable)
  