-----------------------------------------------------------------------------
-- |
-- Module      :  Disco.Desugar
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Desugaring the typechecked surface language to a (still typed)
-- simpler language.
--
-----------------------------------------------------------------------------

module Disco.Desugar
       ( -- * Running desugaring computations
         runDesugar

         -- * Programs, terms, and properties
       , desugarDefn, desugarTerm, desugarProperty

         -- * Case expressions and patterns
       , desugarBranch, desugarGuards
       )
       where

import           Control.Monad.Cont
import           Data.Bool                               (bool)
import           Data.Coerce
import           Data.Maybe                              (fromMaybe, isJust)

import           Disco.AST.Desugared
import           Disco.AST.Surface
import           Disco.AST.Typed
import           Disco.Module
import           Disco.Names
import           Disco.Syntax.Operators
import           Disco.Syntax.Prims
import           Disco.Typecheck                         (containerTy)
import           Disco.Types

import           Disco.Effects.Fresh
import           Polysemy                                (Member, Sem, run)
import           Unbound.Generics.LocallyNameless        (Bind, Name, bind,
                                                          embed, name2String,
                                                          string2Name, unembed,
                                                          unrebind)
import           Unbound.Generics.LocallyNameless.Unsafe (unsafeUnbind)

------------------------------------------------------------
-- Running desugaring computations
------------------------------------------------------------

-- | Run a desugaring computation.
runDesugar :: Sem '[Fresh] a -> a
runDesugar = run . runFresh1
  -- Using runFresh1 is a bit of a hack; that way we won't
  -- ever pick a name with #0 (which is what is generated by default
  -- by string2Name), hence won't conflict with any existing free
  -- variables which came from the parser.

------------------------------------------------------------
-- ATerm DSL
------------------------------------------------------------

-- A tiny DSL for building certain ATerms, which is helpful for
-- writing desugaring rules.

-- Make a local ATVar.
atVar :: Type -> Name ATerm -> ATerm
atVar ty x = ATVar ty (QName LocalName x)

tapp :: ATerm -> ATerm -> ATerm
tapp t1 t2 = ATApp resTy t1 t2
  where
    resTy = case getType t1 of
      (_ :->: r) -> r
      ty         -> error $ "Impossible! Got non-function type " ++ show ty ++ " in tapp"

mkBin :: Type -> BOp -> ATerm -> ATerm -> ATerm
mkBin resTy bop t1 t2
  = tapp (ATPrim (getType t1 :*: getType t2 :->: resTy) (PrimBOp bop)) (mkPair t1 t2)

mkUn :: Type -> UOp -> ATerm -> ATerm
mkUn resTy uop t = tapp (ATPrim (getType t :->: resTy) (PrimUOp uop)) t

mkPair :: ATerm -> ATerm -> ATerm
mkPair t1 t2 = mkTup [t1,t2]

mkTup :: [ATerm] -> ATerm
mkTup ts = ATTup (foldr1 (:*:) (map getType ts)) ts

tapps :: ATerm -> [ATerm] -> ATerm
tapps t ts = tapp t (mkTup ts)

infixr 2 ||.
(||.) :: ATerm -> ATerm -> ATerm
(||.) = mkBin TyBool Or

infixl 6 -., +.
(-.) :: ATerm -> ATerm -> ATerm
at1 -. at2 = mkBin (getType at1) Sub at1 at2

(+.) :: ATerm -> ATerm -> ATerm
at1 +. at2 = mkBin (getType at1) Add at1 at2

infixl 7 /.
(/.) :: ATerm -> ATerm -> ATerm
at1 /. at2 = mkBin (getType at1) Div at1 at2

infix 4 <., >=.
(<.) :: ATerm -> ATerm -> ATerm
(<.) = mkBin TyBool Lt

(>=.) :: ATerm -> ATerm -> ATerm
(>=.) = mkBin TyBool Geq

(|.) :: ATerm -> ATerm -> ATerm
(|.) = mkBin TyBool Divides

infix 4 ==.
(==.) :: ATerm -> ATerm -> ATerm
(==.) = mkBin TyBool Eq

tnot :: ATerm -> ATerm
tnot = tapp (ATPrim (TyBool :->: TyBool) (PrimUOp Not))

(<==.) :: ATerm -> [AGuard] -> ABranch
t <==. gs = bind (toTelescope gs) t

fls :: ATerm
fls = ATBool TyBool False

tru :: ATerm
tru = ATBool TyBool True

tif :: ATerm -> AGuard
tif t = AGBool (embed t)

ctrNil :: Container -> Type -> ATerm
ctrNil ctr ty = ATContainer (containerTy ctr ty) ctr [] Nothing

ctrSingleton :: Container -> ATerm -> ATerm
ctrSingleton ctr t = ATContainer (containerTy ctr (getType t)) ctr [(t, Nothing)] Nothing

------------------------------------------------------------
-- Making DTerms
------------------------------------------------------------

dtVar :: Type -> Name DTerm -> DTerm
dtVar ty x = DTVar ty (QName LocalName x)

dtapp :: DTerm -> DTerm -> DTerm
dtapp t1 t2 = DTApp resTy t1 t2
  where
    resTy = case getType t1 of
      (_ :->: r) -> r
      ty         -> error $ "Impossible! Got non-function type " ++ show ty ++ " in dtapp"

dtbin :: Type -> Prim -> DTerm -> DTerm -> DTerm
dtbin resTy p dt1 dt2
  = dtapp (DTPrim (getType dt1 :*: getType dt2 :->: resTy) p) (mkDTPair dt1 dt2)

mkDTPair :: DTerm -> DTerm -> DTerm
mkDTPair dt1 dt2 = DTPair (getType dt1 :*: getType dt2) dt1 dt2

------------------------------------------------------------
-- Definition desugaring
------------------------------------------------------------

-- | Desugar a definition (consisting of a collection of pattern
--   clauses with bodies) into a core language term.
desugarDefn :: Member Fresh r => Defn -> Sem r DTerm
desugarDefn (Defn _ patTys bodyTy def) =
  desugarAbs Lam (foldr (:->:) bodyTy patTys) def

------------------------------------------------------------
-- Abstraction desugaring
------------------------------------------------------------

-- | Desugar an abstraction -- that is, a collection of clauses
--   with their corresponding patterns. Definitions are abstractions
--   (which happen to be named), and source-level lambdas are also
--   abstractions (which happen to have only one clause).

desugarAbs :: Member Fresh r => Quantifier -> Type -> [Clause] -> Sem r DTerm
-- Special case for compiling a single lambda with no pattern matching directly to a lambda
desugarAbs Lam ty [cl@(unsafeUnbind -> ([APVar _ _], _))] = do
  (ps, at) <- unbind cl
  d <- desugarTerm at
  return $ DTAbs Lam ty (bind (getVar (head ps)) d)
  where
    getVar (APVar _ x) = coerce x
-- General case
desugarAbs quant overallTy body = do
  clausePairs <- unbindClauses body
  let (pats, bodies) = unzip clausePairs
  let patTys = map getType (head pats)
  let bodyTy = getType (head bodies)

  -- generate dummy variables for lambdas
  args <- zipWithM (\_ i -> fresh (string2Name ("arg" ++ show i))) (head pats) [0 :: Int ..]

  -- Create lambdas and one big case.  Recursively desugar the case to
  -- deal with arithmetic patterns.
  let branches = zipWith (mkBranch (zip args patTys)) bodies pats
  dcase <- desugarTerm $ ATCase bodyTy branches
  return $ mkAbs quant overallTy patTys (coerce args) dcase

  where
    mkBranch :: [(Name ATerm, Type)] -> ATerm -> [APattern] -> ABranch
    mkBranch xs b ps = bind (mkGuards xs ps) b

    mkGuards :: [(Name ATerm, Type)] -> [APattern] -> Telescope AGuard
    mkGuards xs ps = toTelescope $ zipWith AGPat (map (\(x,ty) -> embed (atVar ty x)) xs) ps

    -- To make searches fairer, we lift up directly nested abstractions
    -- with the same quantifier when there's only a single clause. That
    -- way, we generate a chain of abstractions followed by a case, instead
    -- of a bunch of alternating abstractions and cases.
    unbindClauses :: Member Fresh r => [Clause] -> Sem r [([APattern], ATerm)]
    unbindClauses [c] | quant `elem` [All, Ex] = do
      (ps, t) <- liftClause c
      return [(ps, addDbgInfo ps t)]
    unbindClauses cs  = mapM unbind cs

    liftClause :: Member Fresh r => Bind [APattern] ATerm -> Sem r ([APattern], ATerm)
    liftClause c = unbind c >>= \case
      (ps, ATAbs q _ c') | q == quant -> do
        (ps', b) <- liftClause c'
        return (ps ++ ps', b)
      (ps, b) -> return (ps, b)

    -- Wrap a term in a test frame to report the values of all variables
    -- bound in the patterns.
    addDbgInfo :: [APattern] -> ATerm -> ATerm
    addDbgInfo ps t = ATTest (map withName $ concatMap varsBound ps) t
      where withName (n, ty) = (name2String n, ty, n)

------------------------------------------------------------
-- Term desugaring
------------------------------------------------------------

-- | Desugar the application of a "bag from counts" primitive on a
--   list, either a safe or unsafe variant.
desugarCList2B :: Member Fresh r => Prim -> Type -> Type -> Type -> Sem r DTerm
desugarCList2B p ty cts b = do
  c <- fresh (string2Name "c")
  body <- desugarTerm $
    tapp (ATPrim (TyBag cts :->: TyBag b) p)
      (tapp (ATPrim (TyList cts :->: TyBag cts) PrimBag)
        (atVar (TyList cts) c)
      )
  return $ mkLambda ty [c] body

-- | Desugar a typechecked term.
desugarTerm :: Member Fresh r => ATerm -> Sem r DTerm
desugarTerm (ATVar ty x) = return $ DTVar ty (coerce x)
desugarTerm (ATPrim (ty1 :->: resTy) (PrimUOp uop))
  | uopDesugars ty1 resTy uop = desugarPrimUOp ty1 resTy uop
desugarTerm (ATPrim (ty1 :*: ty2 :->: resTy) (PrimBOp bop))
  | bopDesugars ty1 ty2 resTy bop = desugarPrimBOp ty1 ty2 resTy bop
desugarTerm (ATPrim ty@(TyList cts :->: TyBag b) PrimC2B) = desugarCList2B PrimC2B ty cts b
desugarTerm (ATPrim ty@(TyList cts :->: TyBag b) PrimUC2B) = desugarCList2B PrimUC2B ty cts b

desugarTerm (ATPrim ty x)        = return $ DTPrim ty x
desugarTerm ATUnit               = return DTUnit
desugarTerm (ATBool ty b)        = return $ DTBool ty b
desugarTerm (ATChar c)           = return $ DTChar c
desugarTerm (ATString cs)        =
  desugarContainer (TyList TyC) ListContainer (map (\c -> (ATChar c, Nothing)) cs) Nothing
desugarTerm (ATAbs q ty lam)     = desugarAbs q ty [lam]

-- Special cases for fully applied operators
desugarTerm (ATApp resTy (ATPrim _ (PrimUOp uop)) t)
  | uopDesugars (getType t) resTy uop = desugarUnApp resTy uop t
desugarTerm (ATApp resTy (ATPrim _ (PrimBOp bop)) (ATTup _ [t1,t2]))
  | bopDesugars (getType t1) (getType t2) resTy bop = desugarBinApp resTy bop t1 t2

desugarTerm (ATApp ty t1 t2)     =
  DTApp ty <$> desugarTerm t1 <*> desugarTerm t2
desugarTerm (ATTup ty ts)        = desugarTuples ty ts
desugarTerm (ATNat ty n)         = return $ DTNat ty n
desugarTerm (ATRat r)            = return $ DTRat r

desugarTerm (ATTyOp ty op t)      = return $ DTTyOp ty op t

desugarTerm (ATChain _ t1 links)  = desugarTerm $ expandChain t1 links

desugarTerm (ATContainer ty c es mell) = desugarContainer ty c es mell

desugarTerm (ATContainerComp _ ctr bqt) = do
  (qs, t) <- unbind bqt
  desugarComp ctr t qs

desugarTerm (ATLet _ t) = do
  (bs, t2) <- unbind t
  desugarLet (fromTelescope bs) t2

desugarTerm (ATCase ty bs) = DTCase ty <$> mapM desugarBranch bs

desugarTerm (ATTest info t) = DTTest (coerce info) <$> desugarTerm t

-- | Desugar a property by wrapping its corresponding term in a test
--   frame to catch its exceptions & convert booleans to props.
desugarProperty :: Member Fresh r => AProperty -> Sem r DTerm
desugarProperty p = DTTest [] <$> desugarTerm p

------------------------------------------------------------
-- Desugaring operators
------------------------------------------------------------

-- | Test whether a given unary operator is one that needs to be
--   desugared, given the type of the argument and result.
uopDesugars :: Type -> Type -> UOp -> Bool
-- uopDesugars _ (TyFin _) Neg = True
uopDesugars _ _         uop = uop == Not

desugarPrimUOp :: Member Fresh r => Type -> Type -> UOp -> Sem r DTerm
desugarPrimUOp argTy resTy op = do
  x <- fresh (string2Name "arg")
  body <- desugarUnApp resTy op (atVar argTy x)
  return $ mkLambda (argTy :->: resTy) [x] body

-- | Test whether a given binary operator is one that needs to be
--   desugared, given the two types of the arguments and the type of the result.
bopDesugars :: Type -> Type -> Type -> BOp -> Bool
bopDesugars _   TyN _ Choose = True
-- bopDesugars _   _   (TyFin _) bop | bop `elem` [Add, Mul] = True
bopDesugars _   _   _ bop = bop `elem`
  [ And, Or, Impl, Iff
  , Neq, Gt, Leq, Geq, Min, Max
  , IDiv
  , Sub, SSub
  , Inter, Diff, Union, Subset
  ]

-- | Desugar a primitive binary operator at the given type.
desugarPrimBOp :: Member Fresh r => Type -> Type -> Type -> BOp -> Sem r DTerm
desugarPrimBOp ty1 ty2 resTy op = do
  p <- fresh (string2Name "pair1")
  x <- fresh (string2Name "arg1")
  y <- fresh (string2Name "arg2")
  let argsTy = ty1 :*: ty2
  body <- desugarBinApp resTy op (atVar ty1 x) (atVar ty2 y)
  return $ mkLambda (argsTy :->: resTy) [p] $
    DTCase resTy
    [ bind
        (toTelescope [DGPat (embed (dtVar argsTy (coerce p))) (DPPair argsTy (coerce x) (coerce y))])
        body
    ]

-- | Desugar a saturated application of a unary operator.
--   The first argument is the type of the result.
desugarUnApp :: Member Fresh r => Type -> UOp -> ATerm -> Sem r DTerm

-- Desugar negation on TyFin to a negation on TyZ followed by a mod.
-- See the comments below re: Add and Mul on TyFin.
-- desugarUnApp (TyFin n) Neg t =
--   desugarTerm $ mkBin (TyFin n) Mod (mkUn TyZ Neg t) (ATNat TyN n)

-- XXX This should be turned into a standard library definition.
-- not t ==> {? false if t, true otherwise ?}
desugarUnApp _ Not t = desugarTerm $
  ATCase TyBool
    [ fls <==. [AGBool (embed t)]
    , tru <==. []
    ]

desugarUnApp ty uop t = error $ "Impossible! desugarUnApp " ++ show ty ++ " " ++ show uop ++ " " ++ show t

-- | Desugar a saturated application of a binary operator.
--   The first argument is the type of the result.
desugarBinApp :: Member Fresh r => Type -> BOp -> ATerm -> ATerm -> Sem r DTerm

-- Implies, and, or should all be turned into a standard library
-- definition.  This will require first (1) adding support for
-- modules/a standard library, including (2) the ability to define
-- infix operators.

-- t1 and t2 ==> {? t2 if t1, false otherwise ?}
desugarBinApp _ And t1 t2 = desugarTerm $
  ATCase TyBool
    [ t2  <==. [tif t1]
    , fls <==. []
    ]

-- (t1 implies t2) ==> (not t1 or t2)
desugarBinApp _ Impl t1 t2 = desugarTerm $ tnot t1 ||. t2

-- (t1 iff t2) ==> (t1 == t2)
desugarBinApp _ Iff t1 t2 = desugarTerm $ t1 ==. t2

-- t1 or t2 ==> {? true if t1, t2 otherwise ?})
desugarBinApp _ Or t1 t2 = desugarTerm $
  ATCase TyBool
    [ tru <==. [tif t1]
    , t2  <==. []
    ]

desugarBinApp _ Neq t1 t2 = desugarTerm $ tnot (t1 ==. t2)
desugarBinApp _ Gt  t1 t2 = desugarTerm $ t2 <. t1
desugarBinApp _ Leq t1 t2 = desugarTerm $ tnot (t2 <. t1)
desugarBinApp _ Geq t1 t2 = desugarTerm $ tnot (t1 <. t2)

-- XXX sharing!
desugarBinApp ty Min t1 t2 = desugarTerm $
  ATCase ty
    [ t1 <==. [tif (t1 <. t2)]
    , t2 <==. []
    ]

desugarBinApp ty Max t1 t2 = desugarTerm $
  ATCase ty
    [ t1 <==. [tif (t2 <. t1)]
    , t2 <==. []
    ]

-- t1 // t2 ==> floor (t1 / t2)
desugarBinApp resTy IDiv t1 t2 = desugarTerm $
  ATApp resTy (ATPrim (getType t1 :->: resTy) PrimFloor) (mkBin (getType t1) Div t1 t2)

-- Desugar normal binomial coefficient (n choose k) to a multinomial
-- coefficient with a singleton list, (n choose [k]).
-- Note this will only be called when (getType t2 == TyN); see bopDesugars.
desugarBinApp _ Choose t1 t2
  = desugarTerm $ mkBin TyN Choose t1 (ctrSingleton ListContainer t2)

desugarBinApp ty Sub  t1 t2 = desugarTerm $ mkBin ty Add t1 (mkUn ty Neg t2)
desugarBinApp ty SSub t1 t2 = desugarTerm $
  -- t1 -. t2 ==> {? 0 if t1 < t2, t1 - t2 otherwise ?}
  ATCase ty
    [ ATNat ty 0         <==. [tif (t1 <. t2)]
    , mkBin ty Sub t1 t2 <==. []
      -- NOTE, the above is slightly bogus since the whole point of SSub is
      -- because we can't subtract naturals.  However, this will
      -- immediately desugar to a DTerm.  When we write a linting
      -- typechecker for DTerms we should allow subtraction on TyN!
    ]

-- Addition and multiplication on TyFin just desugar to the operation
-- followed by a call to mod.
-- desugarBinApp (TyFin n) op t1 t2
--   | op `elem` [Add, Mul] = desugarTerm $
--       mkBin (TyFin n) Mod
--         (mkBin TyN op t1 t2)
--         (ATNat TyN n)
    -- Note the typing of this is a bit funny: t1 and t2 presumably
    -- have type (TyFin n), and now we are saying that applying 'op'
    -- to them results in TyN, then applying 'mod' results in a TyFin
    -- n again.  Using TyN as the intermediate result is necessary so
    -- we don't fall into an infinite desugaring loop, and intuitively
    -- makes sense because the idea is that we first do the operation
    -- as a normal operation in "natural land" and then do a mod.
    --
    -- We will have to think carefully about how the linting
    -- typechecker for DTerms should treat TyN and TyFin.  Probably
    -- something like this will work: TyFin is a subtype of TyN, and
    -- TyN can be turned into TyFin with mod.  (We don't want such
    -- typing rules in the surface disco language itself because
    -- implicit coercions from TyFin -> N don't commute with
    -- operations like addition and multiplication, e.g. 3+3 yields 1
    -- if we add them in Z5 and then coerce to Nat, but 6 if we first
    -- coerce both and then add.

-- Intersection, difference, and union all desugar to an application
-- of 'merge' with an appropriate combining operation.
desugarBinApp ty op t1 t2
  | op `elem` [Inter, Diff, Union] =
    desugarTerm $
    tapps (ATPrim ((TyN :*: TyN :->: TyN) :*: ty :*: ty :->: ty) PrimMerge)
      [ ATPrim (TyN :*: TyN :->: TyN) (mergeOp ty op)
      , t1
      , t2
      ]
  where
    mergeOp _         Inter = PrimBOp Min
    mergeOp _         Diff  = PrimBOp SSub
    mergeOp (TySet _) Union = PrimBOp Max
    mergeOp (TyBag _) Union = PrimBOp Add
    mergeOp _         _     = error $ "Impossible! mergeOp " ++ show ty ++ " " ++ show op

-- A ⊆ B  <==>  (A ⊔ B = B)
--   where ⊔ denotes 'merge max'.
--   Note it is NOT union, since this doesn't work for bags.
--   e.g.  bag [1] union bag [1,2] =  bag [1,1,2] /= bag [1,2].
desugarBinApp _ Subset t1 t2 = desugarTerm $
  tapps (ATPrim (ty :*: ty :->: TyBool) (PrimBOp Eq))
  [ tapps (ATPrim ((TyN :*: TyN :->: TyN) :*: ty :*: ty :->: ty) PrimMerge)
    [ ATPrim (TyN :*: TyN :->: TyN) (PrimBOp Max)
    , t1
    , t2
    ]
  , t2   -- XXX sharing
  ]
  where
    ty = getType t1

desugarBinApp ty bop t1 t2 = error $ "Impossible! desugarBinApp " ++ show ty ++ " " ++ show bop ++ " " ++ show t1 ++ " " ++ show t2

------------------------------------------------------------
-- Desugaring other stuff
------------------------------------------------------------

-- | Desugar a container comprehension.  First translate it into an
--   expanded ATerm and then recursively desugar that.
desugarComp :: Member Fresh r => Container -> ATerm -> Telescope AQual -> Sem r DTerm
desugarComp ctr t qs = expandComp ctr t qs >>= desugarTerm

-- | Expand a container comprehension into an equivalent ATerm.
expandComp :: Member Fresh r => Container -> ATerm -> Telescope AQual -> Sem r ATerm

-- [ t | ] = [ t ]
expandComp ctr t TelEmpty = return $ ctrSingleton ctr t

-- [ t | q, qs ] = ...
expandComp ctr t (TelCons (unrebind -> (q,qs)))
  = case q of
      -- [ t | x in l, qs ] = join (map (\x -> [t | qs]) l)
      AQBind x (unembed -> lst) -> do
        tqs <- expandComp ctr t qs
        let c      = containerTy ctr
            tTy    = getType t
            xTy    = case getType lst of
                       TyContainer _ e -> e
                       _ -> error "Impossible! Not a container in expandComp"
            joinTy = c (c tTy) :->: c tTy
            mapTy  = (xTy :->: c tTy) :*: c xTy :->: c (c tTy)
        return $ tapp (ATPrim joinTy PrimJoin) $
          tapp
            (ATPrim mapTy PrimEach)
            (mkPair
              (ATAbs Lam (xTy :->: c tTy) (bind [APVar xTy x] tqs))
              lst
            )

      -- [ t | g, qs ] = if g then [ t | qs ] else []
      AQGuard (unembed -> g)    -> do
        tqs <- expandComp ctr t qs
        return $ ATCase (containerTy ctr (getType t))
          [ tqs                    <==. [tif g]
          , ctrNil ctr (getType t) <==. []
          ]

-- | Desugar a let into applications of a chain of nested lambdas.
--   /e.g./
--
--     @let x = s, y = t in q@
--
--   desugars to
--
--     @(\x. (\y. q) t) s@
desugarLet :: Member Fresh r => [ABinding] -> ATerm -> Sem r DTerm
desugarLet [] t = desugarTerm t
desugarLet ((ABinding _ x (unembed -> t1)) : ls) t =
  dtapp
    <$> (DTAbs Lam (getType t1 :->: getType t)
          <$> (bind (coerce x) <$> desugarLet ls t)
        )
    <*> desugarTerm t1

-- | Desugar a lambda from a list of argument names and types and the
--   desugared @DTerm@ expression for its body. It will be desugared
--   to a chain of one-argument lambdas. /e.g./
--
--     @\x y z. q@
--
--   desugars to
--
--     @\x. \y. \z. q@
mkLambda :: Type -> [Name ATerm] -> DTerm -> DTerm
mkLambda funty args c = go funty args
  where
    go _ []                    = c
    go ty@(_ :->: ty2) (x:xs) = DTAbs Lam ty (bind (coerce x) (go ty2 xs))
    go ty as = error $ "Impossible! mkLambda.go " ++ show ty ++ " " ++ show as

mkQuant :: Quantifier -> [Type] -> [Name ATerm] -> DTerm -> DTerm
mkQuant q argtys args c = foldr quantify c (zip args argtys)
 where
   quantify (x, ty) body = DTAbs q ty (bind (coerce x) body)

mkAbs :: Quantifier -> Type -> [Type] -> [Name ATerm] -> DTerm -> DTerm
mkAbs Lam funty _ args c = mkLambda funty args c
mkAbs q _ argtys args c  = mkQuant q argtys args c

-- | Desugar a tuple to nested pairs, /e.g./ @(a,b,c,d) ==> (a,(b,(c,d)))@.a
desugarTuples :: Member Fresh r => Type -> [ATerm] -> Sem r DTerm
desugarTuples _ [t]                    = desugarTerm t
desugarTuples ty@(_ :*: ty2) (t:ts) = DTPair ty <$> desugarTerm t <*> desugarTuples ty2 ts
desugarTuples ty ats
  = error $ "Impossible! desugarTuples " ++ show ty ++ " " ++ show ats

-- | Expand a chain of comparisons into a sequence of binary
--   comparisons combined with @and@.  Note we only expand it into
--   another 'ATerm' (which will be recursively desugared), because
--   @and@ itself also gets desugared.
--
--   For example, @a < b <= c > d@ becomes @a < b and b <= c and c > d@.
expandChain :: ATerm -> [ALink] -> ATerm
expandChain _ [] = error "Can't happen! expandChain _ []"
expandChain t1 [ATLink op t2] = mkBin TyBool op t1 t2
expandChain t1 (ATLink op t2 : links) =
  mkBin TyBool And
    (mkBin TyBool op t1 t2)
    (expandChain t2 links)

-- | Desugar a branch of a case expression.
desugarBranch :: Member Fresh r => ABranch -> Sem r DBranch
desugarBranch b = do
  (ags, at) <- unbind b
  dgs <- desugarGuards ags
  d   <- desugarTerm at
  return $ bind dgs d

-- | Desugar the list of guards in one branch of a case expression.
--   Pattern guards essentially remain as they are; boolean guards get
--   turned into pattern guards which match against @true@.
desugarGuards :: Member Fresh r => Telescope AGuard -> Sem r (Telescope DGuard)
desugarGuards = fmap (toTelescope . concat) . mapM desugarGuard . fromTelescope
  where
    desugarGuard :: Member Fresh r => AGuard -> Sem r [DGuard]

    -- A Boolean guard is desugared to a pattern-match on @true = right(unit)@.
    desugarGuard (AGBool (unembed -> at)) = do
      dt <- desugarTerm at
      desugarMatch dt (APInj TyBool R APUnit)

    -- 'let x = t' is desugared to 'when t is x'.
    desugarGuard (AGLet (ABinding _ x (unembed -> at))) = do
      dt <- desugarTerm at
      varMatch dt (coerce x)

    -- Desugaring 'when t is p' is the most complex case; we have to
    -- break down the pattern and match it incrementally.
    desugarGuard (AGPat (unembed -> at) p) = do
      dt <- desugarTerm at
      desugarMatch dt p

    -- Desugar a guard of the form 'when dt is p'.  An entire match is
    -- the right unit to desugar --- as opposed to, say, writing a
    -- function to desugar a pattern --- since a match may desugar to
    -- multiple matches, and on recursive calls we need to know what
    -- term/variable should be bound to the pattern.
    --
    -- A match may desugar to multiple matches for two reasons:
    --
    --   1. Nested patterns 'explode' into a 'telescope' matching one
    --   constructor at a time, for example, 'when t is (x,y,3)'
    --   becomes 'when t is (x,x0) when x0 is (y,x1) when x1 is 3'.
    --   This makes the order of matching explicit and enables lazy
    --   matching without requiring special support from the
    --   interpreter other than WHNF reduction.
    --
    --   2. Matches against arithmetic patterns desugar to a
    --   combination of matching, computation, and boolean checks.
    --   For example, 'when t is (y+1)' becomes 'when t is x0 if x0 >=
    --   1 let y = x0-1'.
    desugarMatch :: Member Fresh r => DTerm -> APattern -> Sem r [DGuard]
    desugarMatch dt (APVar ty x)      = mkMatch dt (DPVar ty (coerce x))
    desugarMatch _  (APWild _)        = return []
    desugarMatch dt APUnit            = mkMatch dt DPUnit
    desugarMatch dt (APBool b)        = desugarMatch dt (APInj TyBool (bool L R b) APUnit)
    desugarMatch dt (APNat ty n)      = desugarMatch (dtbin TyBool (PrimBOp Eq) dt (DTNat ty n)) (APBool True)
    desugarMatch dt (APChar c)        = desugarMatch (dtbin TyBool (PrimBOp Eq) dt (DTChar c)) (APBool True)
    desugarMatch dt (APString s)      = desugarMatch dt (APList (TyList TyC) (map APChar s))
    desugarMatch dt (APTup tupTy pat) = desugarTuplePats tupTy dt pat
      where
        desugarTuplePats :: Member Fresh r => Type -> DTerm -> [APattern] -> Sem r [DGuard]
        desugarTuplePats _ _  [] = error "Impossible! desugarTuplePats []"
        desugarTuplePats _ t [p] = desugarMatch t p
        desugarTuplePats ty@(_ :*: ty2) t (p:ps) = do
          (x1,gs1) <- varForPat p
          (x2,gs2) <- case ps of
            [APVar _ px2] -> return (coerce px2, [])
            _             -> do
              x <- fresh (string2Name "x")
              (x,) <$> desugarTuplePats ty2 (dtVar ty2 x) ps
          fmap concat . sequence $
            [ mkMatch t $ DPPair ty x1 x2
            , return gs1
            , return gs2
            ]
        desugarTuplePats ty _ _
          = error $ "Impossible! desugarTuplePats with non-pair type " ++ show ty

    desugarMatch dt (APInj ty s p) = do
      (x,gs) <- varForPat p
      fmap concat . sequence $
        [ mkMatch dt $ DPInj ty s x
        , return gs
        ]

    desugarMatch dt (APCons ty p1 p2) = do
      y <- fresh (string2Name "y")
      (x1, gs1) <- varForPat p1
      (x2, gs2) <- varForPat p2

      let eltTy = getType p1
          unrolledTy = eltTy :*: ty
      fmap concat . sequence $
        [ mkMatch dt (DPInj ty R y)
        , mkMatch (dtVar unrolledTy y) (DPPair unrolledTy x1 x2)
        , return gs1
        , return gs2
        ]

    desugarMatch dt (APList ty []) = desugarMatch dt (APInj ty L APUnit)
    desugarMatch dt (APList ty ps) =
      desugarMatch dt $ foldr (APCons ty) (APList ty []) ps

    -- when dt is (p + t) ==> when dt is x0; let v = t; [if x0 >= v]; when x0-v is p
    desugarMatch dt (APAdd ty _ p t) = arithBinMatch posRestrict (-.) dt ty p t
      where
        posRestrict plusty
          | plusty `elem` [TyN, TyF] = Just (>=.)
          | otherwise                = Nothing

    -- when dt is (p * t) ==> when dt is x0; let v = t; [if v divides x0]; when x0 / v is p
    desugarMatch dt (APMul ty _ p t) = arithBinMatch intRestrict (/.) dt ty p t
      where
        intRestrict plusty
          | plusty `elem` [TyN, TyZ] = Just (flip (|.))
          | otherwise                = Nothing

    -- when dt is (p - t) ==> when dt is x0; let v = t; when x0 + v is p
    desugarMatch dt (APSub ty p t)  = arithBinMatch (const Nothing) (+.) dt ty p t

    -- when dt is (p/q) ==> when $frac(dt) is (p, q)
    desugarMatch dt (APFrac _ p q)
      = desugarMatch
          (dtapp (DTPrim (TyQ :->: TyZ :*: TyN) PrimFrac) dt)
          (APTup (TyZ :*: TyN) [p, q])

    -- when dt is (-p) ==> when dt is x0; if x0 < 0; when -x0 is p
    desugarMatch dt (APNeg ty p) = do

      -- when dt is x0
      (x0, g1) <- varFor dt

      -- if x0 < 0
      g2  <- desugarGuard $ AGBool (embed (atVar ty (coerce x0) <. ATNat ty 0))

      -- when -x0 is p
      neg <- desugarTerm $ mkUn ty Neg (atVar ty (coerce x0))
      g3  <- desugarMatch neg p

      return (g1 ++ g2 ++ g3)

    mkMatch :: Member Fresh r => DTerm -> DPattern -> Sem r [DGuard]
    mkMatch dt dp = return [DGPat (embed dt) dp]

    varMatch :: Member Fresh r => DTerm -> Name DTerm -> Sem r [DGuard]
    varMatch dt x = mkMatch dt (DPVar (getType dt) x)

    varFor :: Member Fresh r => DTerm -> Sem r (Name DTerm, [DGuard])
    varFor (DTVar _ (QName _ x)) = return (x, [])  -- XXX return a name + provenance??
    varFor dt          = do
      x <- fresh (string2Name "x")
      g <- varMatch dt x
      return (x, g)

    varForPat :: Member Fresh r => APattern -> Sem r (Name DTerm, [DGuard])
    varForPat (APVar _ x) = return (coerce x, [])
    varForPat p           = do
      x <- fresh (string2Name "px")     -- changing this from x fixed a bug and I don't know why =(
      (x,) <$> desugarMatch (dtVar (getType p) x) p

    arithBinMatch
      :: Member Fresh r
      => (Type -> Maybe (ATerm -> ATerm -> ATerm))
      -> (ATerm -> ATerm -> ATerm)
      -> DTerm -> Type -> APattern -> ATerm -> Sem r [DGuard]
    arithBinMatch restrict inverse dt ty p t = do
      (x0, g1) <- varFor dt

      -- let v = t
      t' <- desugarTerm t
      (v, g2) <- varFor t'

      g3 <- case restrict ty of
        Nothing -> return []

        -- if x0 `cmp` v
        Just cmp ->
          desugarGuard $
            AGBool (embed (atVar ty (coerce x0) `cmp` atVar (getType t) (coerce v)))

      -- when x0 `inverse` v is p
      inv <- desugarTerm (atVar ty (coerce x0) `inverse` atVar (getType t) (coerce v))
      g4  <- desugarMatch inv p

      return (g1 ++ g2 ++ g3 ++ g4)

-- | Desugar a container literal such as @[1,2,3]@ or @{1,2,3}@.
desugarContainer :: Member Fresh r => Type -> Container -> [(ATerm, Maybe ATerm)] -> Maybe (Ellipsis ATerm) -> Sem r DTerm

-- Literal list containers desugar to nested applications of cons.
desugarContainer ty ListContainer es Nothing =
  foldr (dtbin ty (PrimBOp Cons)) (DTNil ty) <$> mapM (desugarTerm . fst) es

-- A list container with an ellipsis @[x, y, z .. e]@ desugars to an
-- application of the primitive 'until' function.
desugarContainer ty@(TyList _) ListContainer es (Just (Until t)) =
  dtbin ty PrimUntil
    <$> desugarTerm t
    <*> desugarContainer ty ListContainer es Nothing

-- If desugaring a bag and there are any counts specified, desugar to
-- an application of bagFromCounts to a bag of pairs (with a literal
-- value of 1 filled in for missing counts as needed).
desugarContainer (TyBag eltTy) BagContainer es mell
  | any (isJust . snd) es =
    dtapp (DTPrim (TySet (eltTy :*: TyN) :->: TyBag eltTy) PrimC2B)
      <$> desugarContainer (TyBag (eltTy :*: TyN)) BagContainer counts mell

    where
      -- turn e.g.  x # 3, y   into   (x, 3), (y, 1)
      counts = [ (ATTup (eltTy :*: TyN) [t, fromMaybe (ATNat TyN 1) n], Nothing)
               | (t, n) <- es
               ]

-- Other containers desugar to an application of the appropriate
-- container conversion function to the corresponding desugared list.
desugarContainer ty _ es mell =
  dtapp (DTPrim (TyList eltTy :->: ty) conv)
    <$> desugarContainer (TyList eltTy) ListContainer es mell
  where
    (conv, eltTy) = case ty of
      TyBag e -> (PrimBag, e)
      TySet e -> (PrimSet, e)
      _       -> error $ "Impossible! Non-container type " ++ show ty ++ " in desugarContainer"