content: {-# OPTIONS_GHC -Wall #-}
  
  module Type.Constrain.Expression
    ( constrain,
      constrainDef,
      constrainRecursiveDefs,
    )
  where
  
  import AST.Canonical qualified as Can
  import Data.Index qualified as Index
  import Data.Map.Strict qualified as Map
  import Data.Name qualified as Name
  import Gren.ModuleName qualified as ModuleName
  import Reporting.Annotation qualified as A
  import Reporting.Error.Type (Category (..), Context (..), Expected (..), MaybeName (..), PContext (..), PExpected (..), SubContext (..))
  import Reporting.Error.Type qualified as E
  import Type.Constrain.Pattern qualified as Pattern
  import Type.Instantiate qualified as Instantiate
  import Type.Type as Type hiding (Descriptor (..))
  
  -- CONSTRAIN
  
  -- As we step past type annotations, the free type variables are added to
  -- the "rigid type variables" dict. Allowing sharing of rigid variables
  -- between nested type annotations.
  --
  -- So if you have a top-level type annotation like (func : a -> b) the RTV
  -- dictionary will hold variables for `a` and `b`
  --
  type RTV =
    Map.Map Name.Name Type
  
  constrain :: RTV -> Can.Expr -> Expected Type -> IO Constraint
  constrain rtv (A.At region expression) expected =
    case expression of
      Can.VarLocal name ->
        return (CLocal region name expected)
      Can.VarTopLevel _ name ->
        return (CLocal region name expected)
      Can.VarKernel _ _ ->
        return CTrue
      Can.VarForeign _ name annotation ->
        return $ CForeign region name annotation expected
      Can.VarCtor _ _ name _ annotation ->
        return $ CForeign region name annotation expected
      Can.VarDebug _ name annotation ->
        return $ CForeign region name annotation expected
      Can.VarOperator op _ _ annotation ->
        return $ CForeign region op annotation expected
      Can.Str _ ->
        return $ CEqual region String Type.string expected
      Can.Chr _ ->
        return $ CEqual region Char Type.char expected
      Can.Int _ ->
        do
          var <- mkFlexNumber
          return $ exists [var] $ CEqual region E.Number (VarN var) expected
      Can.Float _ ->
        return $ CEqual region Float Type.float expected
      Can.Array elements ->
        constrainList rtv region elements expected
      Can.Negate expr ->
        do
          numberVar <- mkFlexNumber
          let numberType = VarN numberVar
          numberCon <- constrain rtv expr (FromContext region Negate numberType)
          let negateCon = CEqual region E.Number numberType expected
          return $ exists [numberVar] $ CAnd [numberCon, negateCon]
      Can.Binop op _ _ annotation leftExpr rightExpr ->
        constrainBinop rtv region op annotation leftExpr rightExpr expected
      Can.Lambda args body ->
        constrainLambda rtv region args body expected
      Can.Call func args ->
        constrainCall rtv region func args expected
      Can.If branches finally ->
        constrainIf rtv region branches finally expected
      Can.Case expr branches ->
        constrainCase rtv region expr branches expected
      Can.Let def body ->
        constrainDef rtv def
          =<< constrain rtv body expected
      Can.LetRec defs body ->
        constrainRecursiveDefs rtv defs
          =<< constrain rtv body expected
      Can.LetDestruct pattern expr body ->
        constrainDestruct rtv region pattern expr
          =<< constrain rtv body expected
      Can.Accessor field ->
        do
          extVar <- mkFlexVar
          fieldVar <- mkFlexVar
          let extType = VarN extVar
          let fieldType = VarN fieldVar
          let recordType = RecordN (Map.singleton field fieldType) extType
          return $
            exists [fieldVar, extVar] $
              CEqual region (Accessor field) (FunN recordType fieldType) expected
      Can.Access expr (A.At accessRegion field) ->
        do
          extVar <- mkFlexVar
          fieldVar <- mkFlexVar
          let extType = VarN extVar
          let fieldType = VarN fieldVar
          let recordType = RecordN (Map.singleton field fieldType) extType
  
          let context = RecordAccess (A.toRegion expr) (getAccessName expr) accessRegion field
          recordCon <- constrain rtv expr (FromContext region context recordType)
  
          return $
            exists [fieldVar, extVar] $
              CAnd
                [ recordCon,
                  CEqual region (Access field) fieldType expected
                ]
      Can.Update expr fields ->
        constrainUpdate rtv region expr fields expected
      Can.Record fields ->
        constrainRecord rtv region fields expected
  
  -- CONSTRAIN LAMBDA
  
  constrainLambda :: RTV -> A.Region -> [Can.Pattern] -> Can.Expr -> Expected Type -> IO Constraint
  constrainLambda rtv region args body expected =
    do
      (Args vars tipe resultType (Pattern.State headers pvars revCons)) <-
        constrainArgs args
  
      bodyCon <-
        constrain rtv body (NoExpectation resultType)
  
      return $
        exists vars $
          CAnd
            [ CLet
                { _rigidVars = [],
                  _flexVars = pvars,
                  _header = headers,
                  _headerCon = CAnd (reverse revCons),
                  _bodyCon = bodyCon
                },
              CEqual region Lambda tipe expected
            ]
  
  -- CONSTRAIN CALL
  
  constrainCall :: RTV -> A.Region -> Can.Expr -> [Can.Expr] -> Expected Type -> IO Constraint
  constrainCall rtv region func@(A.At funcRegion _) args expected =
    do
      let maybeName = getName func
  
      funcVar <- mkFlexVar
      resultVar <- mkFlexVar
      let funcType = VarN funcVar
      let resultType = VarN resultVar
  
      funcCon <- constrain rtv func (NoExpectation funcType)
  
      (argVars, argTypes, argCons) <-
        unzip3 <$> Index.indexedTraverse (constrainArg rtv region maybeName) args
  
      let arityType = foldr FunN resultType argTypes
      let category = CallResult maybeName
  
      return $
        exists (funcVar : resultVar : argVars) $
          CAnd
            [ funcCon,
              CEqual funcRegion category funcType (FromContext region (CallArity maybeName (length args)) arityType),
              CAnd argCons,
              CEqual region category resultType expected
            ]
  
  constrainArg :: RTV -> A.Region -> MaybeName -> Index.ZeroBased -> Can.Expr -> IO (Variable, Type, Constraint)
  constrainArg rtv region maybeName index arg =
    do
      argVar <- mkFlexVar
      let argType = VarN argVar
      argCon <- constrain rtv arg (FromContext region (CallArg maybeName index) argType)
      return (argVar, argType, argCon)
  
  getName :: Can.Expr -> MaybeName
  getName (A.At _ expr) =
    case expr of
      Can.VarLocal name -> FuncName name
      Can.VarTopLevel _ name -> FuncName name
      Can.VarForeign _ name _ -> FuncName name
      Can.VarCtor _ _ name _ _ -> CtorName name
      Can.VarOperator op _ _ _ -> OpName op
      Can.VarKernel _ name -> FuncName name
      _ -> NoName
  
  getAccessName :: Can.Expr -> Maybe Name.Name
  getAccessName (A.At _ expr) =
    case expr of
      Can.VarLocal name -> Just name
      Can.VarTopLevel _ name -> Just name
      Can.VarForeign _ name _ -> Just name
      _ -> Nothing
  
  -- CONSTRAIN BINOP
  
  constrainBinop :: RTV -> A.Region -> Name.Name -> Can.Annotation -> Can.Expr -> Can.Expr -> Expected Type -> IO Constraint
  constrainBinop rtv region op annotation leftExpr rightExpr expected =
    do
      leftVar <- mkFlexVar
      rightVar <- mkFlexVar
      answerVar <- mkFlexVar
      let leftType = VarN leftVar
      let rightType = VarN rightVar
      let answerType = VarN answerVar
      let binopType = leftType ==> rightType ==> answerType
  
      let opCon = CForeign region op annotation (NoExpectation binopType)
  
      leftCon <- constrain rtv leftExpr (FromContext region (OpLeft op) leftType)
      rightCon <- constrain rtv rightExpr (FromContext region (OpRight op) rightType)
  
      return $
        exists [leftVar, rightVar, answerVar] $
          CAnd
            [ opCon,
              leftCon,
              rightCon,
              CEqual region (CallResult (OpName op)) answerType expected
            ]
  
  -- CONSTRAIN LISTS
  
  constrainList :: RTV -> A.Region -> [Can.Expr] -> Expected Type -> IO Constraint
  constrainList rtv region entries expected =
    do
      entryVar <- mkFlexVar
      let entryType = VarN entryVar
      let listType = AppN ModuleName.array Name.array [entryType]
  
      entryCons <-
        Index.indexedTraverse (constrainListEntry rtv region entryType) entries
  
      return $
        exists [entryVar] $
          CAnd
            [ CAnd entryCons,
              CEqual region Array listType expected
            ]
  
  constrainListEntry :: RTV -> A.Region -> Type -> Index.ZeroBased -> Can.Expr -> IO Constraint
  constrainListEntry rtv region tipe index expr =
    constrain rtv expr (FromContext region (ArrayEntry index) tipe)
  
  -- CONSTRAIN IF EXPRESSIONS
  
  constrainIf :: RTV -> A.Region -> [(Can.Expr, Can.Expr)] -> Can.Expr -> Expected Type -> IO Constraint
  constrainIf rtv region branches final expected =
    do
      let boolExpect = FromContext region IfCondition Type.bool
      let (conditions, exprs) = foldr (\(c, e) (cs, es) -> (c : cs, e : es)) ([], [final]) branches
  
      condCons <-
        traverse (\c -> constrain rtv c boolExpect) conditions
  
      case expected of
        FromAnnotation name arity _ tipe ->
          do
            branchCons <- Index.indexedForA exprs $ \index expr ->
              constrain rtv expr (FromAnnotation name arity (TypedIfBranch index) tipe)
            return $
              CAnd
                [ CAnd condCons,
                  CAnd branchCons
                ]
        _ ->
          do
            branchVar <- mkFlexVar
            let branchType = VarN branchVar
  
            branchCons <- Index.indexedForA exprs $ \index expr ->
              constrain rtv expr (FromContext region (IfBranch index) branchType)
  
            return $
              exists [branchVar] $
                CAnd
                  [ CAnd condCons,
                    CAnd branchCons,
                    CEqual region If branchType expected
                  ]
  
  -- CONSTRAIN CASE EXPRESSIONS
  
  constrainCase :: RTV -> A.Region -> Can.Expr -> [Can.CaseBranch] -> Expected Type -> IO Constraint
  constrainCase rtv region expr branches expected =
    do
      ptrnVar <- mkFlexVar
      let ptrnType = VarN ptrnVar
      exprCon <- constrain rtv expr (NoExpectation ptrnType)
  
      case expected of
        FromAnnotation name arity _ tipe ->
          do
            branchCons <- Index.indexedForA branches $ \index branch ->
              constrainCaseBranch
                rtv
                branch
                (PFromContext region (PCaseMatch index) ptrnType)
                (FromAnnotation name arity (TypedCaseBranch index) tipe)
  
            return $ exists [ptrnVar] $ CAnd (exprCon : branchCons)
        _ ->
          do
            branchVar <- mkFlexVar
            let branchType = VarN branchVar
  
            branchCons <- Index.indexedForA branches $ \index branch ->
              constrainCaseBranch
                rtv
                branch
                (PFromContext region (PCaseMatch index) ptrnType)
                (FromContext region (CaseBranch index) branchType)
  
            return $
              exists [ptrnVar, branchVar] $
                CAnd
                  [ exprCon,
                    CAnd branchCons,
                    CEqual region Case branchType expected
                  ]
  
  constrainCaseBranch :: RTV -> Can.CaseBranch -> PExpected Type -> Expected Type -> IO Constraint
  constrainCaseBranch rtv (Can.CaseBranch pattern expr) pExpect bExpect =
    do
      (Pattern.State headers pvars revCons) <-
        Pattern.add pattern pExpect Pattern.emptyState
  
      CLet [] pvars headers (CAnd (reverse revCons))
        <$> constrain rtv expr bExpect
  
  -- CONSTRAIN RECORD
  
  constrainRecord :: RTV -> A.Region -> Map.Map (A.Located Name.Name) Can.Expr -> Expected Type -> IO Constraint
  constrainRecord rtv region fields expected =
    do
      dict <- traverse (constrainField rtv) fields
  
      let getType (_, t, _) = t
      let recordType = RecordN (Map.mapKeys A.toValue (Map.map getType dict)) EmptyRecordN
      let recordCon = CEqual region Record recordType expected
  
      let vars = Map.foldr (\(v, _, _) vs -> v : vs) [] dict
      let cons = Map.foldr (\(_, _, c) cs -> c : cs) [recordCon] dict
  
      return $ exists vars (CAnd cons)
  
  constrainField :: RTV -> Can.Expr -> IO (Variable, Type, Constraint)
  constrainField rtv expr =
    do
      var <- mkFlexVar
      let tipe = VarN var
      con <- constrain rtv expr (NoExpectation tipe)
      return (var, tipe, con)
  
  -- CONSTRAIN RECORD UPDATE
  
  constrainUpdate :: RTV -> A.Region -> Can.Expr -> Map.Map (A.Located Name.Name) Can.FieldUpdate -> Expected Type -> IO Constraint
  constrainUpdate rtv region expr locatedFields expected =
    do
      extVar <- mkFlexVar
      let fields = Map.mapKeys A.toValue locatedFields
      fieldDict <- Map.traverseWithKey (constrainUpdateField rtv region) fields
  
      recordVar <- mkFlexVar
      let recordType = VarN recordVar
      let fieldsType = RecordN (Map.map (\(_, t, _) -> t) fieldDict) (VarN extVar)
  
      -- NOTE: fieldsType is separate so that Error propagates better
      let fieldsCon = CEqual region Record recordType (NoExpectation fieldsType)
      let recordCon = CEqual region Record recordType expected
  
      let vars = Map.foldr (\(v, _, _) vs -> v : vs) [recordVar, extVar] fieldDict
      let cons = Map.foldr (\(_, _, c) cs -> c : cs) [recordCon] fieldDict
  
      con <- constrain rtv expr (FromContext region (RecordUpdateKeys fields) recordType)
  
      return $ exists vars $ CAnd (fieldsCon : con : cons)
  
  constrainUpdateField :: RTV -> A.Region -> Name.Name -> Can.FieldUpdate -> IO (Variable, Type, Constraint)
  constrainUpdateField rtv region field (Can.FieldUpdate _ expr) =
    do
      var <- mkFlexVar
      let tipe = VarN var
      con <- constrain rtv expr (FromContext region (RecordUpdateValue field) tipe)
      return (var, tipe, con)
  
  -- CONSTRAIN DESTRUCTURES
  
  constrainDestruct :: RTV -> A.Region -> Can.Pattern -> Can.Expr -> Constraint -> IO Constraint
  constrainDestruct rtv region pattern expr bodyCon =
    do
      patternVar <- mkFlexVar
      let patternType = VarN patternVar
  
      (Pattern.State headers pvars revCons) <-
        Pattern.add pattern (PNoExpectation patternType) Pattern.emptyState
  
      exprCon <-
        constrain rtv expr (FromContext region Destructure patternType)
  
      return $ CLet [] (patternVar : pvars) headers (CAnd (reverse (exprCon : revCons))) bodyCon
  
  -- CONSTRAIN DEF
  
  constrainDef :: RTV -> Can.Def -> Constraint -> IO Constraint
  constrainDef rtv def bodyCon =
    case def of
      Can.Def (A.At region name) args expr ->
        do
          (Args vars tipe resultType (Pattern.State headers pvars revCons)) <-
            constrainArgs args
  
          exprCon <-
            constrain rtv expr (NoExpectation resultType)
  
          return $
            CLet
              { _rigidVars = [],
                _flexVars = vars,
                _header = Map.singleton name (A.At region tipe),
                _headerCon =
                  CLet
                    { _rigidVars = [],
                      _flexVars = pvars,
                      _header = headers,
                      _headerCon = CAnd (reverse revCons),
                      _bodyCon = exprCon
                    },
                _bodyCon = bodyCon
              }
      Can.TypedDef (A.At region name) freeVars typedArgs expr srcResultType ->
        do
          let newNames = Map.difference freeVars rtv
          newRigids <- Map.traverseWithKey (\n _ -> nameToRigid n) newNames
          let newRtv = Map.union rtv (Map.map VarN newRigids)
  
          (TypedArgs tipe resultType (Pattern.State headers pvars revCons)) <-
            constrainTypedArgs newRtv name typedArgs srcResultType
  
          let expected = FromAnnotation name (length typedArgs) TypedBody resultType
          exprCon <-
            constrain newRtv expr expected
  
          return $
            CLet
              { _rigidVars = Map.elems newRigids,
                _flexVars = [],
                _header = Map.singleton name (A.At region tipe),
                _headerCon =
                  CLet
                    { _rigidVars = [],
                      _flexVars = pvars,
                      _header = headers,
                      _headerCon = CAnd (reverse revCons),
                      _bodyCon = exprCon
                    },
                _bodyCon = bodyCon
              }
  
  -- CONSTRAIN RECURSIVE DEFS
  
  data Info = Info
    { _vars :: [Variable],
      _cons :: [Constraint],
      _headers :: Map.Map Name.Name (A.Located Type)
    }
  
  emptyInfo :: Info
  emptyInfo =
    Info [] [] Map.empty
  
  constrainRecursiveDefs :: RTV -> [Can.Def] -> Constraint -> IO Constraint
  constrainRecursiveDefs rtv defs bodyCon =
    recDefsHelp rtv defs bodyCon emptyInfo emptyInfo
  
  recDefsHelp :: RTV -> [Can.Def] -> Constraint -> Info -> Info -> IO Constraint
  recDefsHelp rtv defs bodyCon rigidInfo flexInfo =
    case defs of
      [] ->
        do
          let (Info rigidVars rigidCons rigidHeaders) = rigidInfo
          let (Info flexVars flexCons flexHeaders) = flexInfo
          return $
            CLet rigidVars [] rigidHeaders CTrue $
              CLet [] flexVars flexHeaders (CLet [] [] flexHeaders CTrue (CAnd flexCons)) $
                CAnd [CAnd rigidCons, bodyCon]
      def : otherDefs ->
        case def of
          Can.Def (A.At region name) args expr ->
            do
              let (Info flexVars flexCons flexHeaders) = flexInfo
  
              (Args newFlexVars tipe resultType (Pattern.State headers pvars revCons)) <-
                argsHelp args (Pattern.State Map.empty flexVars [])
  
              exprCon <-
                constrain rtv expr (NoExpectation resultType)
  
              let defCon =
                    CLet
                      { _rigidVars = [],
                        _flexVars = pvars,
                        _header = headers,
                        _headerCon = CAnd (reverse revCons),
                        _bodyCon = exprCon
                      }
  
              recDefsHelp rtv otherDefs bodyCon rigidInfo $
                Info
                  { _vars = newFlexVars,
                    _cons = defCon : flexCons,
                    _headers = Map.insert name (A.At region tipe) flexHeaders
                  }
          Can.TypedDef (A.At region name) freeVars typedArgs expr srcResultType ->
            do
              let newNames = Map.difference freeVars rtv
              newRigids <- Map.traverseWithKey (\n _ -> nameToRigid n) newNames
              let newRtv = Map.union rtv (Map.map VarN newRigids)
  
              (TypedArgs tipe resultType (Pattern.State headers pvars revCons)) <-
                constrainTypedArgs newRtv name typedArgs srcResultType
  
              exprCon <-
                constrain newRtv expr $
                  FromAnnotation name (length typedArgs) TypedBody resultType
  
              let defCon =
                    CLet
                      { _rigidVars = [],
                        _flexVars = pvars,
                        _header = headers,
                        _headerCon = CAnd (reverse revCons),
                        _bodyCon = exprCon
                      }
  
              let (Info rigidVars rigidCons rigidHeaders) = rigidInfo
              recDefsHelp
                rtv
                otherDefs
                bodyCon
                ( Info
                    { _vars = Map.foldr (:) rigidVars newRigids,
                      _cons = CLet (Map.elems newRigids) [] Map.empty defCon CTrue : rigidCons,
                      _headers = Map.insert name (A.At region tipe) rigidHeaders
                    }
                )
                flexInfo
  
  -- CONSTRAIN ARGS
  
  data Args = Args
    { _a_vars :: [Variable],
      _a_type :: Type,
      _a_result :: Type,
      _a_state :: Pattern.State
    }
  
  constrainArgs :: [Can.Pattern] -> IO Args
  constrainArgs args =
    argsHelp args Pattern.emptyState
  
  argsHelp :: [Can.Pattern] -> Pattern.State -> IO Args
  argsHelp args state =
    case args of
      [] ->
        do
          resultVar <- mkFlexVar
          let resultType = VarN resultVar
          return $ Args [resultVar] resultType resultType state
      pattern : otherArgs ->
        do
          argVar <- mkFlexVar
          let argType = VarN argVar
  
          (Args vars tipe result newState) <-
            argsHelp otherArgs
              =<< Pattern.add pattern (PNoExpectation argType) state
  
          return (Args (argVar : vars) (FunN argType tipe) result newState)
  
  -- CONSTRAIN TYPED ARGS
  
  data TypedArgs = TypedArgs
    { _t_type :: Type,
      _t_result :: Type,
      _t_state :: Pattern.State
    }
  
  constrainTypedArgs :: Map.Map Name.Name Type -> Name.Name -> [(Can.Pattern, Can.Type)] -> Can.Type -> IO TypedArgs
  constrainTypedArgs rtv name args srcResultType =
    typedArgsHelp rtv name Index.first args srcResultType Pattern.emptyState
  
  typedArgsHelp :: Map.Map Name.Name Type -> Name.Name -> Index.ZeroBased -> [(Can.Pattern, Can.Type)] -> Can.Type -> Pattern.State -> IO TypedArgs
  typedArgsHelp rtv name index args srcResultType state =
    case args of
      [] ->
        do
          resultType <- Instantiate.fromSrcType rtv srcResultType
          return $ TypedArgs resultType resultType state
      (pattern@(A.At region _), srcType) : otherArgs ->
        do
          argType <- Instantiate.fromSrcType rtv srcType
          let expected = PFromContext region (PTypedArg name index) argType
  
          (TypedArgs tipe resultType newState) <-
            typedArgsHelp rtv name (Index.next index) otherArgs srcResultType
              =<< Pattern.add pattern expected state
  
          return (TypedArgs (FunN argType tipe) resultType newState)
  