{-# LANGUAGE DeriveAnyClass    #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell   #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Disco.Typecheck.Solve
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Constraint solver for the constraints generated during type
-- checking/inference.
-----------------------------------------------------------------------------

module Disco.Typecheck.Solve where

import           Unbound.Generics.LocallyNameless (Alpha, Name, Subst, fv,
                                                   name2Integer, string2Name,
                                                   substs)

import           Data.Coerce
import           GHC.Generics                     (Generic)

import           Control.Arrow                    ((&&&), (***))
import           Control.Lens                     hiding (use, (%=), (.=))
import           Control.Monad                    (unless, zipWithM)
import           Data.Bifunctor                   (first, second)
import           Data.Either                      (partitionEithers)
import           Data.List                        (find, foldl', intersect,
                                                   partition)
import           Data.Map                         (Map, (!))
import qualified Data.Map                         as M
import           Data.Maybe                       (fromJust, fromMaybe,
                                                   mapMaybe)
import           Data.Monoid                      (First (..))
import           Data.Set                         (Set)
import qualified Data.Set                         as S
import           Data.Tuple

import           Disco.Effects.Fresh
import           Disco.Effects.State
import           Polysemy
import           Polysemy.Error
import           Polysemy.Input
import           Polysemy.Output

import           Disco.Messages
import           Disco.Pretty                     hiding ((<>))
import           Disco.Subst
import qualified Disco.Subst                      as Subst
import           Disco.Typecheck.Constraints
import           Disco.Typecheck.Graph            (Graph)
import qualified Disco.Typecheck.Graph            as G
import           Disco.Typecheck.Unify
import           Disco.Types
import           Disco.Types.Qualifiers
import           Disco.Types.Rules

--------------------------------------------------
-- Solver errors

-- | Type of errors which can be generated by the constraint solving
--   process.
data SolveError where
  NoWeakUnifier :: SolveError
  NoUnify       :: SolveError
  UnqualBase    :: Qualifier -> BaseTy    -> SolveError
  Unqual        :: Qualifier -> Type      -> SolveError
  QualSkolem    :: Qualifier -> Name Type -> SolveError
  deriving Show

instance Semigroup SolveError where
  e <> _ = e

--------------------------------------------------
-- Error utilities

runSolve :: Sem (Fresh ': Error SolveError ': r) a -> Sem r (Either SolveError a)
runSolve = runError . runFresh

-- | Run a list of actions, and return the results from those which do
--   not throw an error.  If all of them throw an error, rethrow the
--   first one.
filterErrors :: Member (Error e) r => [Sem r a] -> Sem r [a]
filterErrors ms = do
  es <- mapM try ms
  case partitionEithers es of
    (e:_, []) -> throw e
    (_, as)   -> return as

-- | A variant of 'asum' which picks the first action that succeeds,
--   or re-throws the error of the last one if none of them
--   do. Precondition: the list must not be empty.
asum' :: Member (Error e) r => [Sem r a] -> Sem r a
asum' []     = error "Impossible: asum' []"
asum' [m]    = m
asum' (m:ms) = m `catch` (\_ -> asum' ms)

--------------------------------------------------
-- Simple constraints

data SimpleConstraint where
  (:<:) :: Type -> Type -> SimpleConstraint
  (:=:) :: Type -> Type -> SimpleConstraint
  deriving (Show, Eq, Ord, Generic, Alpha, Subst Type)

instance Pretty SimpleConstraint where
  pretty = \case
    ty1 :<: ty2 -> pretty ty1 <+> "<:" <+> pretty ty2
    ty1 :=: ty2 -> pretty ty1 <+> "=" <+> pretty ty2

--------------------------------------------------
-- Simplifier types

-- Uses TH to generate lenses so it has to go here before other stuff.

---------------------------------
-- Variable maps

-- | Information about a particular type variable.  More information
--   may be added in the future (e.g. polarity).
data TyVarInfo = TVI
  { _tyVarIlk  :: First Ilk   -- ^ The ilk (unification or skolem) of the variable, if known
  , _tyVarSort :: Sort        -- ^ The sort (set of qualifiers) of the type variable.
  }
  deriving (Show)

makeLenses ''TyVarInfo

instance Pretty TyVarInfo where
  pretty (TVI (First ilk) s) = maybe (pure "?") pretty ilk <> "%" <> pretty s

-- | Create a 'TyVarInfo' given an 'Ilk' and a 'Sort'.
mkTVI :: Ilk -> Sort -> TyVarInfo
mkTVI = TVI . First . Just

-- | We can learn different things about a type variable from
--   different places; the 'Semigroup' instance allows combining
--   information about a type variable into a single record.
instance Semigroup TyVarInfo where
  TVI i1 s1 <> TVI i2 s2 = TVI (i1 <> i2) (s1 <> s2)

-- | A 'TyVarInfoMap' records what we know about each type variable;
--   it is a mapping from type variable names to 'TyVarInfo' records.
newtype TyVarInfoMap = VM { unVM :: Map (Name Type) TyVarInfo }
  deriving (Show)

instance Pretty TyVarInfoMap where
  pretty (VM m) = pretty m

-- | Utility function for acting on a 'TyVarInfoMap' by acting on the
--   underlying 'Map'.
onVM ::
  (Map (Name Type) TyVarInfo -> Map (Name Type) TyVarInfo) ->
  TyVarInfoMap -> TyVarInfoMap
onVM f (VM m) = VM (f m)

-- | Look up a given variable name in a 'TyVarInfoMap'.
lookupVM :: Name Type -> TyVarInfoMap -> Maybe TyVarInfo
lookupVM v = M.lookup v . unVM

-- | Remove the mapping for a particular variable name (if it exists)
--   from a 'TyVarInfoMap'.
deleteVM :: Name Type -> TyVarInfoMap -> TyVarInfoMap
deleteVM = onVM . M.delete

-- | Given a list of type variable names, add them all to the
--   'TyVarInfoMap' as 'Skolem' variables (with a trivial sort).
addSkolems :: [Name Type] -> TyVarInfoMap -> TyVarInfoMap
addSkolems vs = onVM $ \vm -> foldl' (flip (\v -> M.insert v (mkTVI Skolem mempty))) vm vs

-- | The @Semigroup@ instance for 'TyVarInfoMap' unions the two maps,
--   combining the info records for any variables occurring in both
--   maps.
instance Semigroup TyVarInfoMap where
  VM sm1 <> VM sm2 = VM (M.unionWith (<>) sm1 sm2)

instance Monoid TyVarInfoMap where
  mempty  = VM M.empty
  mappend = (<>)

-- | Get the sort of a particular variable recorded in a
--   'TyVarInfoMap'.  Returns the trivial (empty) sort for a variable
--   not in the map.
getSort :: TyVarInfoMap -> Name Type -> Sort
getSort (VM m) v = maybe topSort (view tyVarSort) (M.lookup v m)

-- | Get the 'Ilk' of a variable recorded in a 'TyVarInfoMap'.
--   Returns @Nothing@ if the variable is not in the map, or if its
--   ilk is not known.
getIlk :: TyVarInfoMap -> Name Type -> Maybe Ilk
getIlk (VM m) v = (m ^? ix v . tyVarIlk) >>= getFirst

-- | Extend the sort of a type variable by combining its existing sort
--   with the given one.  Has no effect if the variable is not already
--   in the map.
extendSort :: Name Type -> Sort -> TyVarInfoMap -> TyVarInfoMap
extendSort x s = onVM (at x . _Just . tyVarSort %~ (`S.union` s))

---------------------------------
-- Simplifier state

-- The simplification stage maintains a mutable state consisting of
-- the current qualifier map (containing wanted qualifiers for type
-- variables), the list of remaining SimpleConstraints, and the
-- current substitution.  It also keeps track of seen constraints, so
-- expansion of recursive types can stop when encountering a
-- previously seen constraint.
data SimplifyState = SS
  { _ssVarMap      :: TyVarInfoMap
  , _ssConstraints :: [SimpleConstraint]
  , _ssSubst       :: S
  , _ssSeen        :: Set SimpleConstraint
  }

makeLenses ''SimplifyState

lkup :: (Ord k, Show k, Show (Map k a)) => String -> Map k a -> k -> a
lkup messg m k = fromMaybe (error errMsg) (M.lookup k m)
  where
    errMsg = unlines
      [ "Key lookup error:"
      , "  Key = " ++ show k
      , "  Map = " ++ show m
      , "  Location: " ++ messg
      ]

--------------------------------------------------
-- Top-level solver algorithm

solveConstraint
  :: Members '[Fresh, Error SolveError, Output Message, Input TyDefCtx] r
  => Constraint -> Sem r S
solveConstraint c = do

  -- Step 1. Open foralls (instantiating with skolem variables) and
  -- collect wanted qualifiers; also expand disjunctions.  Result in a
  -- list of possible constraint sets; each one consists of equational
  -- and subtyping constraints in addition to qualifiers.

  debug "Solving:"
  debugPretty c

  debug "------------------------------"
  debug "Decomposing constraints..."

  qcList <- decomposeConstraint c

  -- Now try continuing with each set and pick the first one that has
  -- a solution.
  asum' (map (uncurry solveConstraintChoice) qcList)

solveConstraintChoice
  :: Members '[Fresh, Error SolveError, Output Message, Input TyDefCtx] r
  => TyVarInfoMap -> [SimpleConstraint] -> Sem r S
solveConstraintChoice quals cs = do

  debugPretty quals
  debug $ vcat (map pretty' cs)

  tyDefns <- input @TyDefCtx

  -- Step 2. Check for weak unification to ensure termination. (a la
  -- Traytel et al).

  let toEqn (t1 :<: t2) = (t1,t2)
      toEqn (t1 :=: t2) = (t1,t2)
  _ <- note NoWeakUnifier $ weakUnify tyDefns (map toEqn cs)

  -- Step 3. Simplify constraints, resulting in a set of atomic
  -- subtyping constraints.  Also simplify/update qualifier set
  -- accordingly.

  debug "------------------------------"
  debug "Running simplifier..."

  (vm, atoms, theta_simp) <- simplify quals cs
  debug "Done running simplifier. Results:"

  debugPretty vm
  debug $ vcat $ map (pretty' . (\(x,y) -> TyAtom x :<: TyAtom y)) atoms
  debugPretty theta_simp

  -- Step 4. Turn the atomic constraints into a directed constraint
  -- graph.

  debug "------------------------------"
  debug "Generating constraint graph..."

  -- Some variables might have qualifiers but not participate in any
  -- equality or subtyping relations (see issue #153); make sure to
  -- extract them and include them in the constraint graph as isolated
  -- vertices
  let mkAVar (v, First (Just Skolem)) = AVar (S v)
      mkAVar (v, _                  ) = AVar (U v)
      vars = S.fromList . map (mkAVar . second (view tyVarIlk)) . M.assocs . unVM $ vm
      g = mkConstraintGraph vars atoms

  debugPretty g

  -- Step 5.
  -- Check for any weakly connected components containing more
  -- than one skolem, or a skolem and a base type; such components are
  -- not allowed.  Other WCCs with a single skolem simply unify to
  -- that skolem.

  debug "------------------------------"
  debug "Checking WCCs for skolems..."

  (g', theta_skolem) <- checkSkolems vm g
  debugPretty theta_skolem

  -- We don't need to ensure that theta_skolem respects sorts since
  -- checkSkolems will only unify skolem vars with unsorted variables.


  -- Step 6. Eliminate cycles from the graph, turning each strongly
  -- connected component into a single node, unifying all the atoms in
  -- each component.

  debug "------------------------------"
  debug "Collapsing SCCs..."

  (g'', theta_cyc) <- elimCycles tyDefns g'

  debugPretty g''
  debugPretty theta_cyc

  -- Check that the resulting substitution respects sorts...
  let sortOK (x, TyAtom (ABase ty))   = hasSort ty (getSort vm x)
      sortOK (_, TyAtom (AVar (U _))) = True
      sortOK p                        = error $ "Impossible! sortOK " ++ show p
  unless (all sortOK (Subst.toList theta_cyc))
    $ throw NoUnify

  -- ... and update the sort map if we unified any type variables.
  -- Just make sure that if theta_cyc maps x |-> y, then y picks up
  -- the sort of x.

  debug "Old sort map:"
  debugPretty vm

  let vm' = foldr updateVarMap vm (Subst.toList theta_cyc)
        where
          updateVarMap :: (Name Type, Type) -> TyVarInfoMap -> TyVarInfoMap
          updateVarMap (x, TyAtom (AVar (U y))) vmm = extendSort y (getSort vmm x) vmm
          updateVarMap _                        vmm = vmm

  debug "Updated sort map:"
  debugPretty vm

  -- Steps 7 & 8: solve the graph, iteratively finding satisfying
  -- assignments for each type variable based on its successor and
  -- predecessor base types in the graph; then unify all the type
  -- variables in any remaining weakly connected components.

  debug "------------------------------"
  debug "Solving for type variables..."

  theta_sol       <- solveGraph vm' g''
  debugPretty theta_sol

  debug "------------------------------"
  debug "Composing final substitution..."

  let theta_final = theta_sol @@ theta_cyc @@ theta_skolem @@ theta_simp
  debugPretty theta_final

  return theta_final


--------------------------------------------------
-- Step 1. Constraint decomposition.

decomposeConstraint
  :: Members '[Fresh, Error SolveError, Input TyDefCtx] r
  => Constraint -> Sem r [(TyVarInfoMap, [SimpleConstraint])]
decomposeConstraint (CSub t1 t2) = return [(mempty, [t1 :<: t2])]
decomposeConstraint (CEq  t1 t2) = return [(mempty, [t1 :=: t2])]
decomposeConstraint (CQual q ty) = (:[]) . (, []) <$> decomposeQual ty q
decomposeConstraint (CAnd cs)    = map mconcat . sequence <$> mapM decomposeConstraint cs
decomposeConstraint CTrue        = return [mempty]
decomposeConstraint (CAll ty)    = do
  (vars, c) <- unbind ty
  let c' = substs (mkSkolems vars) c
  (map . first . addSkolems) vars <$> decomposeConstraint c'

  where
    mkSkolems :: [Name Type] -> [(Name Type, Type)]
    mkSkolems = map (id &&& TySkolem)

decomposeConstraint (COr cs)     = concat <$> filterErrors (map decomposeConstraint cs)

decomposeQual
  :: Members '[Fresh, Error SolveError, Input TyDefCtx] r
  => Type -> Qualifier -> Sem r TyVarInfoMap
decomposeQual = go S.empty
  where
    go :: Members '[Fresh, Error SolveError, Input TyDefCtx] r
       => Set (String, [Type], Qualifier) -> Type -> Qualifier -> Sem r TyVarInfoMap

    -- For a type atom, call out to checkQual.
    go _ (TyAtom a) q = checkQual q a

    -- Coinductively check user-defined types for a qualifier.  Keep
    -- track of a set of user-defined types and qualifiers we have
    -- seen.  Every time we encounter a new one, add it to the set and
    -- recurse on its unfolding.  If we ever encounter one we have
    -- already seen, we can assume by coinduction that the qualifier
    -- is satisfied.
    go seen (TyCon (CUser t) tys) q = do
      case (t, tys, q) `S.member` seen of
        True -> return mempty
        False -> do
          tyDefns <- input @TyDefCtx
          case M.lookup t tyDefns of
            Nothing  -> error $ show t ++ " not in ty defn map!!"
            Just (TyDefBody _ body) -> do
              let ty' = body tys
              go (S.insert (t, tys, q) seen) ty' q

    -- If we have a container type where the container is still a variable,
    -- just replace it with List for the purposes of generating constraints---
    -- all containers (lists, bags, sets) have the same qualifier rules.
    go seen (TyCon (CContainer (AVar _)) tys) q = go seen (TyCon CList tys) q

    -- Otherwise, decompose a type constructor according to the qualRules.
    go seen ty@(TyCon c tys) q = case qualRules c q of
      Nothing -> throw $ Unqual q ty
      Just qs -> mconcat <$> zipWithM (maybe (return mempty) . go seen) tys qs

checkQual
  :: Members '[Fresh, Error SolveError] r
  => Qualifier -> Atom -> Sem r TyVarInfoMap
checkQual q (AVar (U v)) = return . VM . M.singleton v $ mkTVI Unification (S.singleton q)
checkQual q (AVar (S v)) = throw $ QualSkolem q v
checkQual q (ABase bty)  =
  case hasQual bty q of
    True  -> return mempty
    False -> throw $ UnqualBase q bty

--------------------------------------------------
-- Step 3. Constraint simplification.

-- | This step does unification of equality constraints, as well as
--   structural decomposition of subtyping constraints.  For example,
--   if we have a constraint (x -> y) <: (z -> Int), then we can
--   decompose it into two constraints, (z <: x) and (y <: Int); if we
--   have a constraint v <: (a,b), then we substitute v ↦ (x,y) (where
--   x and y are fresh type variables) and continue; and so on.
--
--   After this step, the remaining constraints will all be atomic
--   constraints, that is, only of the form (v1 <: v2), (v <: b), or
--   (b <: v), where v is a type variable and b is a base type.

simplify
  :: Members '[Error SolveError, Output Message, Input TyDefCtx] r
  => TyVarInfoMap -> [SimpleConstraint] -> Sem r (TyVarInfoMap, [(Atom, Atom)], S)
simplify origVM cs
  = (\(SS vm' cs' s' _) -> (vm', map extractAtoms cs', s'))
  -- contFreshMT :: Monad m => FreshMT m a -> Integer -> m a
  -- "Run a FreshMT computation given a starting index for fresh name generation."
  <$> runFresh' n (execState (SS origVM cs idS S.empty) simplify')
  where

    fvNums :: Alpha a => [a] -> [Integer]
    fvNums = map (name2Integer :: Name Type -> Integer) . toListOf fv

    -- Find first unused integer in constraint free vars and sort map
    -- domain, and use it to start the fresh var generation, so we don't
    -- generate any "fresh" names that interfere with existing names
    n1 = maximum0 . fvNums $ cs
    n = succ . maximum . (n1:) . fvNums . M.keys . unVM $ origVM

    maximum0 [] = 0
    maximum0 xs = maximum xs

    -- Extract the type atoms from an atomic constraint.
    extractAtoms :: SimpleConstraint -> (Atom, Atom)
    extractAtoms (TyAtom a1 :<: TyAtom a2) = (a1, a2)
    extractAtoms c = error $ "Impossible: simplify left non-atomic or non-subtype constraint " ++ show c

    -- Iterate picking one simplifiable constraint and simplifying it
    -- until none are left.
    simplify'
      :: Members '[State SimplifyState, Fresh, Error SolveError, Output Message, Input TyDefCtx] r
      => Sem r ()
    simplify' = do
      -- q <- gets fst
      -- debug (pretty q)
      -- debug ""

      mc <- pickSimplifiable
      case mc of
        Nothing -> return ()
        Just s  -> do

          debug $ "Simplifying:" <+> pretty' s

          simplifyOne s
          simplify'

    -- Pick out one simplifiable constraint, removing it from the list
    -- of constraints in the state.  Return Nothing if no more
    -- constraints can be simplified.
    pickSimplifiable
      :: Members '[State SimplifyState, Fresh, Error SolveError] r
      => Sem r (Maybe SimpleConstraint)
    pickSimplifiable = do
      curCs <- use ssConstraints
      case pick simplifiable curCs of
        Nothing     -> return Nothing
        Just (a,as) -> do
          ssConstraints .= as
          return (Just a)

    -- Pick the first element from a list satisfying the given
    -- predicate, returning the element and the list with the element
    -- removed.
    pick :: (a -> Bool) -> [a] -> Maybe (a,[a])
    pick _ [] = Nothing
    pick p (a:as)
      | p a       = Just (a,as)
      | otherwise = second (a:) <$> pick p as

    -- Check if a constraint can be simplified.  An equality
    -- constraint can always be "simplified" via unification.  A
    -- subtyping constraint can be simplified if either it involves a
    -- type constructor (in which case we can decompose it), or if it
    -- involves two base types (in which case it can be removed if the
    -- relationship holds).
    simplifiable :: SimpleConstraint -> Bool
    simplifiable (_ :=: _)                               = True
    simplifiable (TyCon {} :<: TyCon {})                 = True
    simplifiable (TyVar {} :<: TyCon {})                 = True
    simplifiable (TyCon {} :<: TyVar {})                 = True
    simplifiable (TyCon (CUser _) _ :<: _)               = True
    simplifiable (_ :<: TyCon (CUser _) _)               = True
    simplifiable (TyAtom (ABase _) :<: TyAtom (ABase _)) = True

    simplifiable _                                       = False

    -- Simplify the given simplifiable constraint.  If the constraint
    -- has already been seen before (due to expansion of a recursive
    -- type), just throw it away and stop.
    simplifyOne
      :: Members '[State SimplifyState, Fresh, Error SolveError, Input TyDefCtx] r
      => SimpleConstraint -> Sem r ()
    simplifyOne c = do
      seen <- use ssSeen
      case c `S.member` seen of
        True  -> return ()
        False -> do
          ssSeen %= S.insert c
          simplifyOne' c

    simplifyOne'
      :: Members '[State SimplifyState, Fresh, Error SolveError, Input TyDefCtx] r
      => SimpleConstraint -> Sem r ()

    -- If we have an equality constraint, run unification on it.  The
    -- resulting substitution is applied to the remaining constraints
    -- as well as prepended to the current substitution.

    simplifyOne' (ty1 :=: ty2) = do
      tyDefns <- input @TyDefCtx
      case unify tyDefns [(ty1, ty2)] of
        Nothing -> throw NoUnify
        Just s' -> extendSubst s'

    -- If we see a constraint of the form (T <: a), where T is a
    -- user-defined type and a is a type variable, then just turn it
    -- into an equality (T = a).  This is sound but probably not
    -- complete.  The alternative seems quite complicated, possibly
    -- even undecidable.  See https://github.com/disco-lang/disco/issues/207 .
    simplifyOne' (ty1@(TyCon (CUser _) _) :<: ty2@TyVar{})
      = simplifyOne' (ty1 :=: ty2)

    -- Otherwise, expand the user-defined type and continue.
    simplifyOne' (TyCon (CUser t) ts :<: ty2) = do
      tyDefns <- input @TyDefCtx
      case M.lookup t tyDefns of
        Nothing  -> error $ show t ++ " not in ty defn map!"
        Just (TyDefBody _ body) ->
          ssConstraints %= ((body ts :<: ty2) :)

    -- Turn  a <: T  into  a = T.  See comment above.
    simplifyOne' (ty1@TyVar{} :<: ty2@(TyCon (CUser _) _))
      = simplifyOne' (ty1 :=: ty2)

    simplifyOne' (ty1 :<: TyCon (CUser t) ts) = do
      tyDefns <- input @TyDefCtx
      case M.lookup t tyDefns of
        Nothing  -> error $ show t ++ " not in ty defn map!"
        Just (TyDefBody _ body) ->
          ssConstraints %= ((ty1 :<: body ts) :)

    -- Given a subtyping constraint between two type constructors,
    -- decompose it if the constructors are the same (or fail if they
    -- aren't), taking into account the variance of each argument to
    -- the constructor.  Container types are a special case;
    -- recursively generate a subtyping constraint for their
    -- constructors as well.
    simplifyOne' (TyCon c1@(CContainer ctr1) tys1 :<: TyCon (CContainer ctr2) tys2) =
      ssConstraints %=
        (( (TyAtom ctr1 :<: TyAtom ctr2)
         : zipWith3 variance (arity c1) tys1 tys2
         )
         ++)

    simplifyOne' (TyCon c1 tys1 :<: TyCon c2 tys2)
      | c1 /= c2  = throw NoUnify
      | otherwise =
          ssConstraints %= (zipWith3 variance (arity c1) tys1 tys2 ++)

    -- Given a subtyping constraint between a variable and a type
    -- constructor, expand the variable into the same constructor
    -- applied to fresh type variables.
    simplifyOne' con@(TyVar a   :<: TyCon c _) = expandStruct a c con
    simplifyOne' con@(TyCon c _ :<: TyVar a  ) = expandStruct a c con

    -- Given a subtyping constraint between two base types, just check
    -- whether the first is indeed a subtype of the second.  (Note
    -- that we only pattern match here on type atoms, which could
    -- include variables, but this will only ever get called if
    -- 'simplifiable' was true, which checks that both are base
    -- types.)
    simplifyOne' (TyAtom (ABase b1) :<: TyAtom (ABase b2)) = do
      case isSubB b1 b2 of
        True  -> return ()
        False -> throw NoUnify

    simplifyOne' (s :<: t) =
      error $ "Impossible! simplifyOne' " ++ show s ++ " :<: " ++ show t

    expandStruct
      :: Members '[State SimplifyState, Fresh, Error SolveError, Input TyDefCtx] r
      => Name Type -> Con -> SimpleConstraint -> Sem r ()
    expandStruct a c con = do
      as <- mapM (const (TyVar <$> fresh (string2Name "a"))) (arity c)
      let s' = a |-> TyCon c as
      ssConstraints %= (con:)
      extendSubst s'

    -- 1. compose s' with current subst
    -- 2. apply s' to constraints
    -- 3. apply s' to qualifier map and decompose
    extendSubst
      :: Members '[State SimplifyState, Fresh, Error SolveError, Input TyDefCtx] r
      => S -> Sem r ()
    extendSubst s' = do
      ssSubst %= (s'@@)
      ssConstraints %= applySubst s'
      substVarMap s'

    substVarMap
      :: Members '[State SimplifyState, Fresh, Error SolveError, Input TyDefCtx] r
      => S -> Sem r ()
    substVarMap s' = do
      vm <- use ssVarMap

      -- 1. Get quals for each var in domain of s' and match them with
      -- the types being substituted for those vars.

      let tySorts :: [(Type, Sort)]
          tySorts = map (second (view tyVarSort)) . mapMaybe (traverse (`lookupVM` vm) . swap) $ Subst.toList s'

          tyQualList :: [(Type, Qualifier)]
          tyQualList = concatMap (sequenceA . second S.toList) tySorts

      -- 2. Decompose the resulting qualifier constraints

      vm' <- mconcat <$> mapM (uncurry decomposeQual) tyQualList

      -- 3. delete domain of s' from vm and merge in decomposed quals.

      ssVarMap .= vm' <> foldl' (flip deleteVM) vm (dom s')

      -- The above works even when unifying two variables.  Say we have
      -- the TyVarInfoMap
      --
      --   a |-> {add}
      --   b |-> {sub}
      --
      -- and we get back theta = [a |-> b].  The domain of theta
      -- consists solely of a, so we look up a in the TyVarInfoMap and get
      -- {add}.  We therefore generate the constraint 'add (theta a)'
      -- = 'add b' which can't be decomposed at all, and hence yields
      -- the TyVarInfoMap {b |-> {add}}.  We then delete a from the
      -- original TyVarInfoMap and merge the result with the new TyVarInfoMap,
      -- yielding {b |-> {sub,add}}.


    -- Create a subtyping constraint based on the variance of a type
    -- constructor argument position: in the usual order for
    -- covariant, and reversed for contravariant.
    variance Co     ty1 ty2 = ty1 :<: ty2
    variance Contra ty1 ty2 = ty2 :<: ty1

--------------------------------------------------
-- Step 4: Build constraint graph

-- | Given a list of atoms and atomic subtype constraints (each pair
--   @(a1,a2)@ corresponds to the constraint @a1 <: a2@) build the
--   corresponding constraint graph.
mkConstraintGraph :: (Show a, Ord a) => Set a -> [(a, a)] -> Graph a
mkConstraintGraph as cs = G.mkGraph nodes (S.fromList cs)
  where
    nodes = as `S.union` S.fromList (cs ^.. traverse . each)

--------------------------------------------------
-- Step 5: Check skolems

-- | Check for any weakly connected components containing more than
--   one skolem, or a skolem and a base type, or a skolem and any
--   variables with nontrivial sorts; such components are not allowed.
--   If there are any WCCs with a single skolem, no base types, and
--   only unsorted variables, just unify them all with the skolem and
--   remove those components.
checkSkolems
  :: Members '[Error SolveError, Output Message, Input TyDefCtx] r
  => TyVarInfoMap -> Graph Atom -> Sem r (Graph UAtom, S)
checkSkolems vm graph = do
  let skolemWCCs :: [Set Atom]
      skolemWCCs = filter (any isSkolem) $ G.wcc graph

      ok wcc =  S.size (S.filter isSkolem wcc) <= 1
             && all (\case { ABase _    -> False
                           ; AVar (S _) -> True
                           ; AVar (U v) -> maybe True (S.null . view tyVarSort) (lookupVM v vm) })
                wcc

      (good, bad) = partition ok skolemWCCs

  unless (null bad) $ throw NoUnify

  -- take all good sets and
  --   (1) delete them from the graph
  --   (2) unify them all with the skolem
  unifyWCCs graph idS good

  where
    noSkolems :: Atom -> UAtom
    noSkolems (ABase b)    = UB b
    noSkolems (AVar (U v)) = UV v
    noSkolems (AVar (S v)) = error $ "Skolem " ++ show v ++ " remaining in noSkolems"

    unifyWCCs
      :: Members '[Error SolveError, Output Message, Input TyDefCtx] r
      => Graph Atom -> S -> [Set Atom] -> Sem r (Graph UAtom, S)
    unifyWCCs g s []     = return (G.map noSkolems g, s)
    unifyWCCs g s (u:us) = do
      debug $ "Unifying" <+> pretty' (u:us) <> "..."

      tyDefns <- input @TyDefCtx

      let g' = foldl' (flip G.delete) g u

          ms' = unifyAtoms tyDefns (S.toList u)
      case ms' of
        Nothing -> throw NoUnify
        Just s' -> unifyWCCs g' (atomToTypeSubst s' @@ s) us

--------------------------------------------------
-- Step 6: Eliminate cycles

-- | Eliminate cycles in the constraint set by collapsing each
--   strongly connected component to a single node, (unifying all the
--   types in the SCC). A strongly connected component is a maximal
--   set of nodes where every node is reachable from every other by a
--   directed path; since we are using directed edges to indicate a
--   subtyping constraint, this means every node must be a subtype of
--   every other, and the only way this can happen is if all are in
--   fact equal.
--
--   Of course, this step can fail if the types in a SCC are not
--   unifiable.  If it succeeds, it returns the collapsed graph (which
--   is now guaranteed to be acyclic, i.e. a DAG) and a substitution.
elimCycles
  :: Members '[Error SolveError] r
  => TyDefCtx -> Graph UAtom -> Sem r (Graph UAtom, S)
elimCycles tyDefns = elimCyclesGen uatomToTypeSubst (unifyUAtoms tyDefns)

elimCyclesGen
  :: forall a b r. (Subst a a, Ord a, Members '[Error SolveError] r)
  => (Substitution a -> Substitution b) -> ([a] -> Maybe (Substitution a))
  -> Graph a -> Sem r (Graph a, Substitution b)
elimCyclesGen genSubst genUnify g
  = note NoUnify
  $ (G.map fst &&& (genSubst . compose . S.map snd . G.nodes)) <$> g'
  where

    g' :: Maybe (Graph (a, Substitution a))
    g' = G.sequenceGraph $ G.map unifySCC (G.condensation g)

    unifySCC :: Set a -> Maybe (a, Substitution a)
    unifySCC uatoms = case S.toList uatoms of
      []       -> error "Impossible! unifySCC on the empty set"
      as@(a:_) -> (flip applySubst a &&& id) <$> genUnify as

------------------------------------------------------------
-- Steps 7 and 8: Constraint resolution
------------------------------------------------------------

-- | Rels stores the set of base types and variables related to a
--   given variable in the constraint graph (either predecessors or
--   successors, but not both).
data Rels = Rels
  { baseRels :: Set BaseTy
  , varRels  :: Set (Name Type)
  }
  deriving (Show, Eq)

-- | A RelMap associates each variable to its sets of base type and
--   variable predecessors and successors in the constraint graph.
newtype RelMap = RelMap { unRelMap :: Map (Name Type, Dir) Rels}

instance Pretty RelMap where
  pretty (RelMap rm) = vcat (map prettyVar byVar)
    where
      vars = S.map fst (M.keysSet rm)
      byVar = map (\x -> (rm!(x,SubTy), x, rm!(x,SuperTy))) (S.toList vars)

      prettyVar (subs, x, sups) = hsep [prettyRel subs, "<:", pretty x, "<:", prettyRel sups]
      prettyRel rs = pretty (baseRels rs) <> ", " <> pretty (varRels rs)

-- | Modify a @RelMap@ to record the fact that we have solved for a
--   type variable.  In particular, delete the variable from the
--   @RelMap@ as a key, and also update the relative sets of every
--   other variable to remove this variable and add the base type we
--   chose for it.
substRel :: Name Type -> BaseTy -> RelMap -> RelMap
substRel x ty
  = RelMap
  . M.delete (x,SuperTy)
  . M.delete (x,SubTy)
  . M.map (\r@(Rels b v) -> if x `S.member` v then Rels (S.insert ty b) (S.delete x v) else r)
  . unRelMap

-- | Essentially dirtypesBySort vm rm dir t s x finds all the
--   dir-types (sub- or super-) of t which have sort s, relative to
--   the variables in x.  This is \overbar{T}_S^X (resp. \underbar...)
--   from Traytel et al.
dirtypesBySort :: TyVarInfoMap -> RelMap -> Dir -> BaseTy -> Sort -> Set (Name Type) -> [BaseTy]
dirtypesBySort vm (RelMap relMap) dir t s x

    -- Keep only those supertypes t' of t
  = keep (dirtypes dir t) $ \t' ->
      -- which have the right sort, and such that
      hasSort t' s &&

      -- for all variables beta \in x,
      forAll x (\beta ->

        -- there is at least one type t'' which is a subtype of t'
        -- which would be a valid solution for beta, that is,
        exists (dirtypes (other dir) t') $ \t'' ->

          -- t'' has the sort beta is supposed to have, and
          hasSort t'' (getSort vm beta) &&

          -- t'' is a supertype of every base type predecessor of beta.
          forAll (baseRels (lkup "dirtypesBySort, beta rel" relMap (beta, other dir)))
            (isDirB dir t''))

    -- The above comments are written assuming dir = Super; of course,
    -- if dir = Sub then just swap "super" and "sub" everywhere.

  where
    forAll, exists :: Foldable t => t a -> (a -> Bool) -> Bool
    forAll = flip all
    exists = flip any
    keep   = flip filter

-- | Sort-aware infimum or supremum.
limBySort :: TyVarInfoMap -> RelMap -> Dir -> [BaseTy] -> Sort -> Set (Name Type) -> Maybe BaseTy
limBySort vm rm dir ts s x
  = (\is -> find (\lim -> all (\u -> isDirB dir u lim) is) is)
  . isects
  . map (\t -> dirtypesBySort vm rm dir t s x)
  $ ts
  where
    isects = foldr1 intersect

lubBySort, glbBySort :: TyVarInfoMap -> RelMap -> [BaseTy] -> Sort -> Set (Name Type) -> Maybe BaseTy
lubBySort vm rm = limBySort vm rm SuperTy
glbBySort vm rm = limBySort vm rm SubTy

-- | From the constraint graph, build the sets of sub- and super- base
--   types of each type variable, as well as the sets of sub- and
--   supertype variables.  For each type variable x in turn, try to
--   find a common supertype of its base subtypes which is consistent
--   with the sort of x and with the sorts of all its sub-variables,
--   as well as symmetrically a common subtype of its supertypes, etc.
--   Assign x one of the two: if it has only successors, assign it
--   their inf; otherwise, assign it the sup of its predecessors.  If
--   it has both, we have a choice of whether to assign it the sup of
--   predecessors or inf of successors; both lead to a sound &
--   complete algorithm.  We choose to assign it the sup of its
--   predecessors in this case, since it seems nice to default to
--   "simpler" types lower down in the subtyping chain.
solveGraph
  :: Members '[Fresh, Error SolveError, Output Message] r
  => TyVarInfoMap -> Graph UAtom -> Sem r S
solveGraph vm g = atomToTypeSubst . unifyWCC <$> go topRelMap
  where
    unifyWCC :: Substitution BaseTy -> Substitution Atom
    unifyWCC s = compose (map mkEquateSubst wccVarGroups) @@ fmap ABase s
      where
        wccVarGroups :: [Set (Name Type)]
        wccVarGroups  = map (S.map getVar) . filter (all uisVar) . applySubst s $ G.wcc g
        getVar (UV v) = v
        getVar (UB b) = error
          $ "Impossible! Base type " ++ show b ++ " in solveGraph.getVar"

        mkEquateSubst :: Set (Name Type) -> Substitution Atom
        mkEquateSubst = mkEquations . S.toList

        mkEquations (a:as) = Subst.fromList . map (\v -> (coerce v, AVar (U a))) $ as
        mkEquations []     = error "Impossible! Empty set of names in mkEquateSubst"

            -- After picking concrete base types for all the type
            -- variables we can, the only thing possibly remaining in
            -- the graph are components containing only type variables
            -- and no base types.  It is sound, and simplifies the
            -- generated types considerably, to simply unify any type
            -- variables which are related by subtyping constraints.
            -- That is, we collect all the type variables in each
            -- weakly connected component and unify them.
            --
            -- As an example where this final step makes a difference,
            -- consider a term like @\x. (\y.y) x@.  A fresh type
            -- variable is generated for the type of @x@, and another
            -- for the type of @y@; the application of @(\y.y)@ to @x@
            -- induces a subtyping constraint between the two type
            -- variables.  The most general type would be something
            -- like @forall a b. (a <: b) => a -> b@, but we want to
            -- avoid generating unnecessary subtyping constraints (the
            -- type system might not even support subtyping qualifiers
            -- like this).  Instead, we unify the two type variables
            -- and the resulting type is @forall a. a -> a@.

    -- Get the successor and predecessor sets for all the type variables.
    topRelMap :: RelMap
    topRelMap
      = RelMap
      . M.map (uncurry Rels . (S.fromAscList *** S.fromAscList)
               . partitionEithers . map uatomToEither . S.toList)
      $ M.mapKeys (,SuperTy) subMap `M.union` M.mapKeys (,SubTy) superMap

    subMap, superMap :: Map (Name Type) (Set UAtom)
    (subMap, superMap) = (onlyVars *** onlyVars) $ G.cessors g

    onlyVars :: Map UAtom (Set UAtom) -> Map (Name Type) (Set UAtom)
    onlyVars = M.mapKeys fromVar . M.filterWithKey (\a _ -> uisVar a)
      where
        fromVar (UV x) = x
        fromVar _      = error "Impossible! UB but uisVar."

    go
      :: Members '[Fresh, Error SolveError, Output Message] r
      => RelMap -> Sem r (Substitution BaseTy)
    go relMap@(RelMap rm) = debugPretty relMap >> case as of

      -- No variables left that have base type constraints.
      []    -> return idS

      -- Solve one variable at a time.  See below.
      (a:_) -> do
        debug $ "Solving for" <+> pretty' a
        case solveVar a of
          Nothing       -> do
            debug $ "Couldn't solve for" <+> pretty' a
            throw NoUnify

          -- If we solved for a, delete it from the maps, apply the
          -- resulting substitution to the remainder (updating the
          -- relMap appropriately), and recurse.  The substitution we
          -- want will be the composition of the substitution for a
          -- with the substitution generated by the recursive call.
          --
          -- Note we don't need to delete a from the TyVarInfoMap; we
          -- never use the set of keys from the TyVarInfoMap for
          -- anything (indeed, some variables might not be keys if
          -- they have an empty sort), so it doesn't matter if old
          -- variables hang around in it.
          Just s -> do
            debugPretty s
            (@@ s) <$> go (substRel a (fromJust $ Subst.lookup (coerce a) s) relMap)

      where
        -- NOTE we can't solve a bunch in parallel!  Might end up
        -- assigning them conflicting solutions if some depend on
        -- others.  For example, consider the situation
        --
        --            Z
        --            |
        --            a3
        --           /  \
        --          a1   N
        --
        -- If we try to solve in parallel we will end up assigning a1
        -- -> Z (since it only has base types as an upper bound) and
        -- a3 -> N (since it has both upper and lower bounds, and by
        -- default we pick the lower bound), but this is wrong since
        -- we should have a1 < a3.
        --
        -- If instead we solve them one at a time, we could e.g. first
        -- solve a1 -> Z, and then we would find a3 -> Z as well.
        -- Alternately, if we first solve a3 -> N then we will have a1
        -- -> N as well.  Both are acceptable.
        --
        -- In fact, this exact graph comes from (^x.x+1) which was
        -- erroneously being inferred to have type Z -> N when I first
        -- wrote the code.

        -- Get only the variables we can solve on this pass, which
        -- have base types in their predecessor or successor set.  If
        -- there are no such variables, then start picking any
        -- remaining variables with a sort and pick types for them
        -- (disco doesn't have qualified polymorphism so we can't just
        -- leave them).
        asBase
          = map fst
          . filter (not . S.null . baseRels . lkup "solveGraph.go.as" rm)
          $ M.keys rm
        as = case asBase of
          [] -> filter ((/= topSort) . getSort vm) . map fst $ M.keys rm
          _  -> asBase

        -- Solve for a variable, failing if it has no solution, otherwise returning
        -- a substitution for it.
        solveVar :: Name Type -> Maybe (Substitution BaseTy)
        solveVar v =
          case ((v,SuperTy), (v,SubTy)) & over both (S.toList . baseRels . lkup "solveGraph.solveVar" rm) of
            -- No sub- or supertypes; the only way this can happen is
            -- if it has a nontrivial sort.
            --
            -- Traytel et al. don't seem to have a rule saying what to
            -- do in this case (see Fig. 16 on p. 16 of their long
            -- version).  We used to just pick a type that inhabits
            -- the sort, but this is wrong; see
            -- https://github.com/disco-lang/disco/issues/192.
            --
            -- For now, let's assume that any situation in which we
            -- have no base sub- or supertypes but we do have
            -- nontrivial sorts means that we are dealing with numeric
            -- types; so we can just call N a base subtype and go from there.

            ([], []) ->
              -- Debug.trace (show v ++ " has no sub- or supertypes.  Assuming N as a subtype.")
              (coerce v |->) <$> lubBySort vm relMap [N] (getSort vm v)
                (varRels (lkup "solveVar none, rels" rm (v,SubTy)))

            -- Only supertypes.  Just assign a to their inf, if one exists.
            (bsupers, []) ->
              -- Debug.trace (show v ++ " has only supertypes (" ++ show bsupers ++ ")") $
              (coerce v |->) <$> glbBySort vm relMap bsupers (getSort vm v)
                (varRels (lkup "solveVar bsupers, rels" rm (v,SuperTy)))

            -- Only subtypes.  Just assign a to their sup.
            ([], bsubs)   ->
              -- Debug.trace (show v ++ " has only subtypes (" ++ show bsubs ++ ")") $
              -- Debug.trace ("sortmap: " ++ show vm) $
              -- Debug.trace ("relmap: " ++ show relMap) $
              -- Debug.trace ("sort for " ++ show v ++ ": " ++ show (getSort vm v)) $
              -- Debug.trace ("relvars: " ++ show (varRels (relMap ! (v,SubTy)))) $
              (coerce v |->) <$> lubBySort vm relMap bsubs (getSort vm v)
                (varRels (lkup "solveVar bsubs, rels" rm (v,SubTy)))

            -- Both successors and predecessors.  Both must have a
            -- valid bound, and the bounds must not overlap.  Assign a
            -- to the sup of its predecessors.
            (bsupers, bsubs) -> do
              ub <- glbBySort vm relMap bsupers (getSort vm v)
                      (varRels (rm ! (v,SuperTy)))
              lb <- lubBySort vm relMap bsubs   (getSort vm v)
                      (varRels (rm ! (v,SubTy)))
              case isSubB lb ub of
                True  -> Just (coerce v |-> lb)
                False -> Nothing