module Core.LinearCheck

import Core.CaseTree
import Core.Context
import Core.Context.Log
import Core.Core
import Core.Env
import Core.Normalise
import Core.Options
import Core.UnifyState
import Core.Value
import Core.TT

import Data.List

%default covering

-- List of variable usages - we'll count the contents of specific variables
-- when discharging binders, to ensure that linear names are only used once
data Usage : List Name -> Type where
     Nil : Usage vars
     (::) : Var vars -> Usage vars -> Usage vars

Show (Usage vars) where
  show xs = "[" ++ showAll xs ++ "]"
    where
      showAll : Usage vs -> String
      showAll [] = ""
      showAll [el] = show el
      showAll (x :: xs) = show x ++ ", " ++ show xs

doneScope : Usage (n :: vars) -> Usage vars
doneScope [] = []
doneScope (MkVar First :: xs) = doneScope xs
doneScope (MkVar (Later p) :: xs) = MkVar p :: doneScope xs

(++) : Usage ns -> Usage ns -> Usage ns
(++) [] ys = ys
(++) (x :: xs) ys = x :: xs ++ ys

count : Nat -> Usage ns -> Nat
count p [] = 0
count p (v :: xs)
    = if p == varIdx v then 1 + count p xs else count p xs

localPrf : {later : _} -> Var (later ++ n :: vars)
localPrf {later = []} = MkVar First
localPrf {n} {vars} {later = (x :: xs)}
    = let MkVar p = localPrf {n} {vars} {later = xs} in
          MkVar (Later p)

mutual
  updateHoleUsageArgs : {vars : _} ->
                        {auto c : Ref Ctxt Defs} ->
                        {auto u : Ref UST UState} ->
                        (useInHole : Bool) ->
                        Var vars -> List (Var vars) ->
                        List (Term vars) -> Core Bool
  updateHoleUsageArgs useInHole var zs [] = pure False
  updateHoleUsageArgs useInHole var zs (a :: as)
      = do h <- updateHoleUsage useInHole var zs a
           h' <- updateHoleUsageArgs useInHole var zs as
           pure (h || h')

  -- The assumption here is that hole types are abstracted over the entire
  -- environment, so that they have the appropriate number of function
  -- arguments and there are no lets
  updateHoleType : {vars : _} ->
                   {auto c : Ref Ctxt Defs} ->
                   {auto u : Ref UST UState} ->
                   (useInHole : Bool) ->
                   Var vars -> List (Var vars) ->
                   Term vs -> List (Term vars) ->
                   Core (Term vs)
  updateHoleType useInHole var zs (Bind bfc nm (Pi fc' c e ty) sc) (Local _ r v _ :: as)
      -- if the argument to the hole type is the variable of interest,
      -- and the variable should be used in the hole, set it to Rig1,
      -- otherwise set it to Rig0
      = if varIdx var == v
           then do scty <- updateHoleType False var zs sc as
                   let c' = if useInHole then c else erased
                   pure (Bind bfc nm (Pi fc' c' e ty) scty)
           else if elem v (map varIdx zs)
                then do scty <- updateHoleType useInHole var zs sc as
                        pure (Bind bfc nm (Pi fc' erased e ty) scty)
                else do scty <- updateHoleType useInHole var zs sc as
                        pure (Bind bfc nm (Pi fc' c e ty) scty)
  updateHoleType useInHole var zs (Bind bfc nm (Pi fc' c e ty) sc) (a :: as)
      = do ignore $ updateHoleUsage False var zs a
           scty <- updateHoleType useInHole var zs sc as
           pure (Bind bfc nm (Pi fc' c e ty) scty)
  updateHoleType useInHole var zs ty as
      = do ignore $ updateHoleUsageArgs False var zs as
           pure ty

  updateHoleUsagePats : {auto c : Ref Ctxt Defs} ->
                        {auto u : Ref UST UState} ->
                        (useInHole : Bool) ->
                        Var vars -> List (Term vars) ->
                        (vs ** (Env Term vs, Term vs, Term vs)) ->
                        Core Bool
  updateHoleUsagePats {vars} useInHole var args (vs ** (env, lhs, rhs))
      = do -- Find the argument which corresponds to var
           let argpos = findArg Z args
           log "quantity.hole" 10 $ "At positions " ++ show argpos
           -- Find what it's position is in env by looking at the lhs args
           let vars = mapMaybe (findLocal (getArgs lhs)) argpos
           hs <- traverse (\vsel => updateHoleUsage useInHole vsel [] rhs)
                          vars
           pure (any id hs)
    where
      findArg : Nat -> List (Term vars) -> List Nat
      findArg i [] = []
      findArg i (Local _ _ idx vel :: els)
          = if idx == varIdx var
               then i :: findArg (1 + i) els
               else findArg (1 + i) els
      findArg i (_ :: els) = findArg (1 + i) els

      findLocal : List (Term vs) -> Nat -> Maybe (Var vs)
      findLocal (Local _ _ _ p :: _) Z = Just (MkVar p)
      findLocal (As _ _ (Local _ _ _ p) _ :: _) Z = Just (MkVar p)
      findLocal (As _ _ _ (Local _ _ _ p) :: _) Z = Just (MkVar p)
      findLocal (_ :: els) (S k) = findLocal els k
      findLocal _ _ = Nothing

  updateHoleUsage : {vars : _} ->
                    {auto c : Ref Ctxt Defs} ->
                    {auto u : Ref UST UState} ->
                    (useInHole : Bool) ->
                    Var vars -> List (Var vars) ->
                    Term vars -> Core Bool
  updateHoleUsage useInHole (MkVar var) zs (Bind _ _ (Let _ _ val _) sc)
      = do h <- updateHoleUsage useInHole (MkVar var) zs val
           h' <- updateHoleUsage useInHole (MkVar (Later var)) (map weaken zs) sc
           pure (h || h')
  updateHoleUsage useInHole (MkVar var) zs (Bind _ n b sc)
      = updateHoleUsage useInHole (MkVar (Later var)) (map weaken zs) sc
  updateHoleUsage useInHole var zs (Meta fc n i args)
      = do defs <- get Ctxt
           Just gdef <- lookupCtxtExact (Resolved i) (gamma defs)
                | Nothing => updateHoleUsageArgs useInHole var zs args
           -- only update for holes with no definition yet
           case definition gdef of
                Hole _ _ =>
                   do let ty = type gdef
                      ty' <- updateHoleType useInHole var zs ty args
                      updateTy i ty'
                      logTerm "quantity.hole.update" 5 ("New type of " ++
                                 show (fullname gdef)) ty'
                      logTerm "quantity.hole.update" 5 ("Updated from " ++
                                 show (fullname gdef)) (type gdef)
                      pure True
                _ => updateHoleUsageArgs useInHole var zs args
  updateHoleUsage useInHole var zs (As _ _ a p)
      = do h <- updateHoleUsage useInHole var zs a
           h' <- updateHoleUsage useInHole var zs a
           pure (h || h')
  updateHoleUsage useInHole var zs (TDelayed _ _ t)
      = updateHoleUsage useInHole var zs t
  updateHoleUsage useInHole var zs (TDelay _ _ _ t)
      = updateHoleUsage useInHole var zs t
  updateHoleUsage useInHole var zs (TForce _ _ t)
      = updateHoleUsage useInHole var zs t
  updateHoleUsage useInHole var zs tm
      = case getFnArgs tm of
             (Ref _ _ fn, args) =>
                  -- no need to look inside 'fn' for holes since we did that
                  -- when working through lcheckDef recursively
                  updateHoleUsageArgs useInHole var zs args
             (f, []) => pure False
             (f, args) => updateHoleUsageArgs useInHole var zs (f :: args)

-- Linearity checking of an already checked term. This serves two purposes:
--  + Checking correct usage of linear bindings
--  + updating hole types to reflect usage counts correctly
-- Returns term, normalised type, and a list of used variables
mutual
  lcheck : {vars : _} ->
           {auto c : Ref Ctxt Defs} ->
           {auto u : Ref UST UState } ->
           RigCount -> (erase : Bool) -> Env Term vars -> Term vars ->
           Core (Term vars, Glued vars, Usage vars)
  lcheck {vars} rig erase env (Local {name} fc x idx prf)
      = let b = getBinder prf env
            rigb = multiplicity b
            ty = binderType b in
            do when (not erase) $ rigSafe rigb rig
               pure (Local fc x idx prf, gnf env ty, used rig)
    where
      getName : {idx : _} -> (vs : List Name) -> (0 p : IsVar n idx vs) -> Name
      getName (x :: _) First = x
      getName (x :: xs) (Later p) = getName xs p

      rigSafe : RigCount -> RigCount -> Core ()
      rigSafe l r = when (l < r)
                         (throw (LinearMisuse fc (getName vars prf) l r))

      -- count the usage if we're in a linear context. If not, the usage doesn't
      -- matter
      used : RigCount -> Usage vars
      used r = if isLinear r then [MkVar prf] else []

  lcheck rig erase env (Ref fc nt fn)
      = do ty <- lcheckDef fc rig erase env fn
           pure (Ref fc nt fn, gnf env (embed ty), [])

  -- If the meta has a definition, and we're not in Rig0, expand it first
  -- and check the result.
  -- Otherwise, don't count variable usage in holes, so as far as linearity
  -- checking is concerned, update the type so that the binders
  -- are in Rig0
  lcheck {vars} rig erase env (Meta fc n idx args)
      = do defs <- get Ctxt
           Just gdef <- lookupCtxtExact (Resolved idx) (gamma defs)
                | _ => undefinedName fc n
           let expand = branchZero
                          (case type gdef of
                                Erased _ _ => True -- defined elsewhere, need to expand
                                _ => False)
                          (case definition gdef of
                                (PMDef _ _ _ _ _) => True
                                _ => False)
                          rig
           logC "quantity" 10 $ do
             def <- case definition gdef of
                         PMDef _ _ (STerm _ tm) _ _ =>
                              do tm' <- toFullNames tm
                                 pure (show tm')
                         _ => pure ""
             pure (show rig ++ ": " ++ show n ++ " " ++ show fc ++ "\n"
                     ++ show def)
           if expand
              then expandMeta rig erase env n idx (definition gdef) args
              else do let ty : ClosedTerm
                             = case definition gdef of
                                    Hole _ _ => unusedHoleArgs args (type gdef)
                                    _ => type gdef
                      nty <- nf defs env (embed ty)
                      lcheckMeta rig erase env fc n idx args [] nty
    where
      unusedHoleArgs : List a -> Term vs -> Term vs
      unusedHoleArgs (_ :: args) (Bind bfc n (Pi fc _ e ty) sc)
          = Bind bfc n (Pi fc erased e ty) (unusedHoleArgs args sc)
      unusedHoleArgs args (Bind bfc n (Let fc c e ty) sc)
          = Bind bfc n (Let fc c e ty) (unusedHoleArgs args sc)
      unusedHoleArgs _ ty = ty

  lcheck rig_in erase env (Bind fc nm b sc)
      = do (b', bt, usedb) <- handleUnify (lcheckBinder rig erase env b)
                                 (\err =>
                                     case err of
                                          LinearMisuse _ _ r _ =>
                                             lcheckBinder rig erase env
                                                (setMultiplicity b linear)
                                          _ => throw err)
           -- Anything linear can't be used in the scope of a lambda, if we're
           -- checking in general context
           let env' = if rig_in == top
                         then case b of
                              (Lam _ _ _ _) => eraseLinear env
                              _ => env
                         else env
           (sc', sct, usedsc) <- lcheck rig erase (b' :: env') sc
           defs <- get Ctxt

           let used_in = count 0 usedsc
           holeFound <- if not erase && isLinear (multiplicity b)
                           then updateHoleUsage (used_in == 0)
                                         (MkVar First)
                                         (map weaken (getZeroes env'))
                                         sc'
                           else pure False

           -- if there's a hole, assume it will contain the missing usage
           -- if there is none already
           let used = if isLinear ((multiplicity b) |*| rig) &&
                         holeFound && used_in == 0
                         then 1
                         else used_in

           when (not erase) $
               checkUsageOK used ((multiplicity b) |*| rig)
           defs <- get Ctxt
           discharge defs env fc nm b' bt sc' sct (usedb ++ doneScope usedsc)
    where
      rig : RigCount
      rig = case b of
                 Pi _ _ _ _ =>
                      if isErased rig_in
                         then erased
                         else top -- checking as if an inspectable run-time type
                 Let _ _ _ _ => rig_in
                 _ => if isErased rig_in
                         then erased
                         else linear

      getZeroes : {vs : _} -> Env Term vs -> List (Var vs)
      getZeroes [] = []
      getZeroes (b :: bs)
          = if isErased (multiplicity b)
               then MkVar First :: map weaken (getZeroes bs)
               else map weaken (getZeroes bs)

      eraseLinear : Env Term vs -> Env Term vs
      eraseLinear [] = []
      eraseLinear (b :: bs)
          = if isLinear (multiplicity b)
               then setMultiplicity b erased :: eraseLinear bs
               else b :: eraseLinear bs

      checkUsageOK : Nat -> RigCount -> Core ()
      checkUsageOK used r = when (isLinear r && used /= 1)
                                 (throw (LinearUsed fc used nm))

  lcheck rig erase env (App fc f a)
      = do (f', gfty, fused) <- lcheck rig erase env f
           defs <- get Ctxt
           fty <- getNF gfty
           case fty of
                NBind _ _ (Pi _ rigf _ ty) scdone =>
                     -- if the argument is borrowed, it's okay to use it in
                     -- unrestricted context, because we'll be out of the
                     -- application without spending it
                   do let checkRig = rigf |*| rig
                      (a', gaty, aused) <- lcheck checkRig erase env a
                      sc' <- scdone defs (toClosure defaultOpts env a')
                      let aerased = if erase && isErased rigf then Erased fc False else a'
                      -- Possibly remove this check, or make it a compiler
                      -- flag? It is a useful double check on the result of
                      -- elaboration, but there are pathological cases where
                      -- it makes the check very slow (id id id id ... id id etc
                      -- for example) and there may be similar realistic cases.
                      -- If elaboration is correct, this should never fail!
                      opts <- getSession
                      when (debugElabCheck opts) $ do
                        aty <- getNF gaty
                        when (not !(convert defs env aty !(evalClosure defs ty))) $
                           do ty' <- quote defs env ty
                              aty' <- quote defs env aty
                              throw (CantConvert fc (gamma defs) env ty' aty')
                      pure (App fc f' aerased,
                            glueBack defs env sc',
                            fused ++ aused)
                NApp _ (NRef _ n) _ =>
                      do Just _ <- lookupCtxtExact n (gamma defs)
                              | _ => undefinedName fc n
                         tfty <- getTerm gfty
                         throw (GenericMsg fc ("Linearity checking failed on " ++ show !(toFullNames f') ++
                              " (" ++ show !(toFullNames tfty) ++ " not a function type)"))

                _ => do tfty <- getTerm gfty
                        throw (GenericMsg fc ("Linearity checking failed on " ++ show !(toFullNames f') ++
                              " (" ++ show !(toFullNames tfty) ++ " not a function type)"))

  lcheck rig erase env (As fc s as pat)
      = do (as', _, _) <- lcheck rig erase env as
           (pat', pty, u) <- lcheck rig erase env pat
           pure (As fc s as' pat', pty, u)
  lcheck rig erase env (TDelayed fc r ty)
      = do (ty', _, u) <- lcheck rig erase env ty
           pure (TDelayed fc r ty', gType fc, u)
  lcheck rig erase env (TDelay fc r ty val)
      = do (ty', _, _) <- lcheck erased erase env ty
           (val', gty, u) <- lcheck rig erase env val
           ty <- getTerm gty
           pure (TDelay fc r ty' val', gnf env (TDelayed fc r ty), u)
  lcheck rig erase env (TForce fc r val)
      = do (val', gty, u) <- lcheck rig erase env val
           tynf <- getNF gty
           case tynf of
                NDelayed _ r narg
                    => do defs <- get Ctxt
                          pure (TForce fc r val', glueBack defs env narg, u)
                _ => throw (GenericMsg fc "Not a delayed type")
  lcheck rig erase env (PrimVal fc c)
      = pure (PrimVal fc c, gErased fc, [])
  lcheck rig erase env (Erased fc i)
      = pure (Erased fc i, gErased fc, [])
  lcheck rig erase env (TType fc)
      = pure (TType fc, gType fc, [])

  lcheckBinder : {vars : _} ->
                 {auto c : Ref Ctxt Defs} ->
                 {auto u : Ref UST UState} ->
                 RigCount -> (erase : Bool) -> Env Term vars ->
                 Binder (Term vars) ->
                 Core (Binder (Term vars), Glued vars, Usage vars)
  lcheckBinder rig erase env (Lam fc c x ty)
      = do (tyv, tyt, _) <- lcheck erased erase env ty
           pure (Lam fc c x tyv, tyt, [])
  lcheckBinder rig erase env (Let fc rigc val ty)
      = do (tyv, tyt, _) <- lcheck erased erase env ty
           (valv, valt, vs) <- lcheck (rig |*| rigc) erase env val
           pure (Let fc rigc valv tyv, tyt, vs)
  lcheckBinder rig erase env (Pi fc c x ty)
      = do (tyv, tyt, _) <- lcheck (rig |*| c) erase env ty
           pure (Pi fc c x tyv, tyt, [])
  lcheckBinder rig erase env (PVar fc c p ty)
      = do (tyv, tyt, _) <- lcheck erased erase env ty
           pure (PVar fc c p tyv, tyt, [])
  lcheckBinder rig erase env (PLet fc rigc val ty)
      = do (tyv, tyt, _) <- lcheck erased erase env ty
           (valv, valt, vs) <- lcheck (rig |*| rigc) erase env val
           pure (PLet fc rigc valv tyv, tyt, vs)
  lcheckBinder rig erase env (PVTy fc c ty)
      = do (tyv, tyt, _) <- lcheck erased erase env ty
           pure (PVTy fc c tyv, tyt, [])

  discharge : {vars : _} ->
              Defs -> Env Term vars ->
              FC -> (nm : Name) -> Binder (Term vars) -> Glued vars ->
              Term (nm :: vars) -> Glued (nm :: vars) -> Usage vars ->
              Core (Term vars, Glued vars, Usage vars)
  discharge defs env fc nm (Lam fc' c x ty) gbindty scope gscopety used
       = do scty <- getTerm gscopety
            pure (Bind fc nm (Lam fc' c x ty) scope,
                  gnf env (Bind fc nm (Pi fc' c x ty) scty), used)
  discharge defs env fc nm (Let fc' c val ty) gbindty scope gscopety used
       = do scty <- getTerm gscopety
            pure (Bind fc nm (Let fc' c val ty) scope,
                  gnf env (Bind fc nm (Let fc' c val ty) scty), used)
  discharge defs env fc nm (Pi fc' c x ty) gbindty scope gscopety used
       = pure (Bind fc nm (Pi fc' c x ty) scope, gbindty, used)
  discharge defs env fc nm (PVar fc' c p ty) gbindty scope gscopety used
       = do scty <- getTerm gscopety
            pure (Bind fc nm (PVar fc' c p ty) scope,
                  gnf env (Bind fc nm (PVTy fc' c ty) scty), used)
  discharge defs env fc nm (PLet fc' c val ty) gbindty scope gscopety used
       = do scty <- getTerm gscopety
            pure (Bind fc nm (PLet fc' c val ty) scope,
                  gnf env (Bind fc nm (PLet fc' c val ty) scty), used)
  discharge defs env fc nm (PVTy fc' c ty) gbindty scope gscopety used
       = pure (Bind fc nm (PVTy fc' c ty) scope, gbindty, used)

  data ArgUsage
       = UseAny -- RigW so we don't care
       | Use0 -- argument position not used
       | Use1 -- argument position used exactly once
       | UseKeep -- keep as is
       | UseUnknown -- hole, so can't tell

  Show ArgUsage where
    show UseAny = "any"
    show Use0 = "0"
    show Use1 = "1"
    show UseKeep = "keep"
    show UseUnknown = "unknown"

  -- Check argument usage in case blocks. Returns a list of how each argument
  -- in the case block is used, to build the appropriate type for the outer
  -- block.
  getArgUsage : {auto c : Ref Ctxt Defs} ->
                {auto e : Ref UST UState} ->
                FC -> RigCount -> ClosedTerm ->
                List (vs ** (Env Term vs, Term vs, Term vs)) ->
                Core (List ArgUsage)
  getArgUsage topfc rig ty pats
      = do us <- traverse (getPUsage ty) pats
           pure (map snd !(combine us))
    where
      getCaseUsage : {vs : _} ->
                     Term ns -> Env Term vs -> List (Term vs) ->
                     Usage vs -> Term vs ->
                     Core (List (Name, ArgUsage))
      getCaseUsage ty env (As _ _ _ p :: args) used rhs
          = getCaseUsage ty env (p :: args) used rhs
      getCaseUsage (Bind _ n (Pi _ rig _ ty) sc) env (arg :: args) used rhs
          = if isLinear rig
               then case arg of
                         (Local _ _ idx p) =>
                           do rest <- getCaseUsage sc env args used rhs
                              let used_in = count idx used
                              holeFound <- updateHoleUsage (used_in == 0) (MkVar p) [] rhs
                              let ause
                                  = if holeFound && used_in == 0
                                            then UseUnknown
                                            else if used_in == 0
                                                    then Use0
                                                    else Use1
                              pure ((n, ause) :: rest)
                         _ => do elseCase
               else elseCase
          where
            elseCase : Core (List (Name, ArgUsage))
            elseCase = do rest <- getCaseUsage sc env args used rhs
                          pure $ if isErased rig
                             then ((n, Use0) :: rest)
                             else ((n, UseKeep) :: rest)
      getCaseUsage tm env args used rhs = pure []

      checkUsageOK : FC -> Nat -> Name -> Bool -> RigCount -> Core ()
      checkUsageOK fc used nm isloc rig
          = when (isLinear rig && ((isloc && used > 1) || (not isloc && used /= 1)))
                 (throw (LinearUsed fc used nm))

      -- Is the variable one of the lhs arguments; i.e. do we treat it as
      -- affine rather than linear
      isLocArg : Var vars -> List (Term vars) -> Bool
      isLocArg p [] = False
      isLocArg p (Local _ _ idx _ :: args)
          = idx == varIdx p || isLocArg p args
      isLocArg p (As _ _ tm pat :: args)
          = isLocArg p (tm :: pat :: args)
      isLocArg p (_ :: args) = isLocArg p args

      -- As checkEnvUsage in general, but it's okay for local variables to
      -- remain unused (since in that case, they must be used outside the
      -- case block)
      checkEnvUsage : {done, vars : _} ->
                      RigCount ->
                      Env Term vars -> Usage (done ++ vars) ->
                      List (Term (done ++ vars)) ->
                      Term (done ++ vars) -> Core ()
      checkEnvUsage rig [] usage args tm = pure ()
      checkEnvUsage rig {done} {vars = nm :: xs} (b :: env) usage args tm
          = do let pos = localPrf {later = done}
               let used_in = count (varIdx pos) usage

               holeFound <- if isLinear (multiplicity b)
                               then updateHoleUsage (used_in == 0) pos [] tm
                               else pure False
               let used = if isLinear ((multiplicity b) |*| rig) &&
                             holeFound && used_in == 0
                             then 1
                             else used_in
               checkUsageOK (getLoc (binderType b))
                            used nm (isLocArg pos args)
                                    ((multiplicity b) |*| rig)
               checkEnvUsage {done = done ++ [nm]} rig env
                     (rewrite sym (appendAssociative done [nm] xs) in usage)
                     (rewrite sym (appendAssociative done [nm] xs) in args)
                     (rewrite sym (appendAssociative done [nm] xs) in tm)

      getPUsage : ClosedTerm -> (vs ** (Env Term vs, Term vs, Term vs)) ->
                  Core (List (Name, ArgUsage))
      getPUsage ty (_ ** (penv, lhs, rhs))
          = do logEnv "quantity" 10 "Env" penv
               logTerm "quantity" 10 "LHS" lhs
               logTerm "quantity" 5 "Linear check in case RHS" rhs
               (rhs', _, used) <- lcheck rig False penv rhs
               log "quantity" 10 $ "Used: " ++ show used
               let args = getArgs lhs
               checkEnvUsage {done = []} rig penv used args rhs'
               ause <- getCaseUsage ty penv args used rhs
               log "quantity" 10 $ "Arg usage: " ++ show ause
               pure ause

      combineUsage : (Name, ArgUsage) -> (Name, ArgUsage) ->
                     Core (Name, ArgUsage)
      combineUsage (n, Use0) (_, Use1)
          = throw (GenericMsg topfc ("Inconsistent usage of " ++ show n ++ " in case branches"))
      combineUsage (n, Use1) (_, Use0)
          = throw (GenericMsg topfc ("Inconsistent usage of " ++ show n ++ " in case branches"))
      combineUsage (n, UseAny) _ = pure (n, UseAny)
      combineUsage _ (n, UseAny) = pure (n, UseAny)
      combineUsage (n, UseKeep) _ = pure (n, UseKeep)
      combineUsage _ (n, UseKeep) = pure (n, UseKeep)
      combineUsage (n, UseUnknown) _ = pure (n, UseUnknown)
      combineUsage _ (n, UseUnknown) = pure (n, UseUnknown)
      combineUsage x y = pure x

      combineUsages : List (Name, ArgUsage) -> List (Name, ArgUsage) ->
                      Core (List (Name, ArgUsage))
      combineUsages [] [] = pure []
      combineUsages (u :: us) (v :: vs)
          = do u' <- combineUsage u v
               us' <- combineUsages us vs
               pure (u' :: us')
      combineUsages _ _ = throw (InternalError "Argument usage lists inconsistent")

      combine : List (List (Name, ArgUsage)) ->
                Core (List (Name, ArgUsage))
      combine [] = pure []
      combine [x] = pure x
      combine (x :: xs)
          = do xs' <- combine xs
               combineUsages x xs'

  lcheckDef : {auto c : Ref Ctxt Defs} ->
              {auto u : Ref UST UState} ->
              FC -> RigCount -> (erase : Bool) -> Env Term vars -> Name ->
              Core ClosedTerm
  lcheckDef fc rig True env n
      = do defs <- get Ctxt
           Just def <- lookupCtxtExact n (gamma defs)
                | Nothing => undefinedName fc n
           pure (type def)
  lcheckDef fc rig False env n
      = do defs <- get Ctxt
           let Just idx = getNameID n (gamma defs)
                | Nothing => undefinedName fc n
           Just def <- lookupCtxtExact (Resolved idx) (gamma defs)
                | Nothing => undefinedName fc n
           rigSafe (multiplicity def) rig
           if linearChecked def
              then pure (type def)
              else case definition def of
                        PMDef _ _ _ _ pats =>
                            do u <- getArgUsage (getLoc (type def))
                                                rig (type def) pats
                               log "quantity" 5 $ "Overall arg usage " ++ show u
                               let ty' = updateUsage u (type def)
                               updateTy idx ty'
                               setLinearCheck idx True
                               logTerm "quantity" 5 ("New type of " ++
                                          show (fullname def)) ty'
                               logTerm "quantity" 5 ("Updated from " ++
                                          show (fullname def)) (type def)
                               pure ty'
                        _ => pure (type def)
    where
      updateUsage : List ArgUsage -> Term ns -> Term ns
      updateUsage (u :: us) (Bind bfc n (Pi fc c e ty) sc)
          = let sc' = updateUsage us sc
                c' = case u of
                          Use0 => erased
                          Use1 => linear -- ignore usage elsewhere, we checked here
                          UseUnknown => c -- don't know, assumed unchanged and update hole types
                          UseKeep => c -- matched here, so count usage elsewhere
                          UseAny => c in -- no constraint, so leave alone
                Bind bfc n (Pi fc c' e ty) sc'
      updateUsage _ ty = ty

      rigSafe : RigCount -> RigCount -> Core ()
      rigSafe a b = when (a < b)
                         (throw (LinearMisuse fc !(getFullName n) a b))

  expandMeta : {vars : _} ->
               {auto c : Ref Ctxt Defs} ->
               {auto u : Ref UST UState} ->
               RigCount -> (erase : Bool) -> Env Term vars ->
               Name -> Int -> Def -> List (Term vars) ->
               Core (Term vars, Glued vars, Usage vars)
  expandMeta rig erase env n idx (PMDef _ [] (STerm _ fn) _ _) args
      = do tm <- substMeta (embed fn) args []
           lcheck rig erase env tm
    where
      substMeta : {drop, vs : _} ->
                  Term (drop ++ vs) -> List (Term vs) -> SubstEnv drop vs ->
                  Core (Term vs)
      substMeta (Bind bfc n (Lam _ c e ty) sc) (a :: as) env
          = substMeta sc as (a :: env)
      substMeta (Bind bfc n (Let _ c val ty) sc) as env
          = substMeta (subst val sc) as env
      substMeta rhs [] env = pure (substs env rhs)
      substMeta rhs _ _ = throw (InternalError ("Badly formed metavar solution " ++ show n ++ " " ++ show fn))
  expandMeta rig erase env n idx def _
      = throw (InternalError ("Badly formed metavar solution " ++ show n ++ " " ++ show def))

  lcheckMeta : {vars : _} ->
               {auto c : Ref Ctxt Defs} ->
               {auto u : Ref UST UState} ->
               RigCount -> Bool -> Env Term vars ->
               FC -> Name -> Int ->
               (args : List (Term vars)) ->
               (checked : List (Term vars)) ->
               NF vars -> Core (Term vars, Glued vars, Usage vars)
  lcheckMeta rig erase env fc n idx
             (arg :: args) chk (NBind _ _ (Pi _ rigf _ ty) sc)
      = do let checkRig = rigf |*| rig
           (arg', gargTy, aused) <- lcheck checkRig erase env arg
           defs <- get Ctxt
           sc' <- sc defs (toClosure defaultOpts env arg')
           let aerased = if erase && isErased rigf then Erased fc False else arg'
           (tm, gty, u) <- lcheckMeta rig erase env fc n idx args
                                      (aerased :: chk) sc'
           pure (tm, gty, aused ++ u)
  lcheckMeta rig erase env fc n idx (arg :: args) chk nty
      = do defs <- get Ctxt
           empty <- clearDefs defs
           ty <- quote empty env nty
           throw (GenericMsg fc ("Linearity checking failed on metavar "
                      ++ show !(toFullNames n) ++ " (" ++ show !(toFullNames ty)
                      ++ " not a function type)"))
  lcheckMeta rig erase env fc n idx [] chk nty
      = do defs <- get Ctxt
           pure (Meta fc n idx (reverse chk), glueBack defs env nty, [])

checkEnvUsage : {vars, done : _} ->
                {auto c : Ref Ctxt Defs} ->
                {auto u : Ref UST UState} ->
                FC -> RigCount ->
                Env Term vars -> Usage (done ++ vars) ->
                Term (done ++ vars) ->
                Core ()
checkEnvUsage fc rig [] usage tm = pure ()
checkEnvUsage fc rig {done} {vars = nm :: xs} (b :: env) usage tm
    = do let pos = localPrf {later = done}
         let used_in = count (varIdx pos) usage

         holeFound <- if isLinear (multiplicity b)
                         then updateHoleUsage (used_in == 0) pos [] tm
                         else pure False
         let used = if isLinear ((multiplicity b) |*| rig) &&
                       holeFound && used_in == 0
                       then 1
                       else used_in
         checkUsageOK used ((multiplicity b) |*| rig)
         checkEnvUsage {done = done ++ [nm]} fc rig env
               (rewrite sym (appendAssociative done [nm] xs) in usage)
               (rewrite sym (appendAssociative done [nm] xs) in tm)
  where
    checkUsageOK : Nat -> RigCount -> Core ()
    checkUsageOK used r = when (isLinear r && used /= 1)
                               (throw (LinearUsed fc used nm))

-- Linearity check an elaborated term. If 'erase' is set, erase anything that's in
-- a Rig0 argument position (we can't do this until typechecking is complete, though,
-- since it might be used for unification/reasoning elsewhere, so we only do this for
-- definitions ready for compilation).
export
linearCheck : {vars : _} ->
              {auto c : Ref Ctxt Defs} ->
              {auto u : Ref UST UState} ->
              FC -> RigCount -> (erase : Bool) ->
              Env Term vars -> Term vars ->
              Core (Term vars)
linearCheck fc rig erase env tm
    = do logTerm "quantity" 5 "Linearity check on " tm
         logEnv "quantity" 5 "In env" env
         (tm', _, used) <- lcheck rig erase env tm
         log "quantity" 5 $ "Used: " ++ show used
         when (not erase) $ checkEnvUsage {done = []} fc rig env used tm'
         pure tm'


||| Copyright 2016 Google Inc.
|||
||| Licensed under the Apache License, Version 2.0 (the "License");
||| you may not use this file except in compliance with the License.
||| You may obtain a copy of the License at
|||
|||     http://www.apache.org/licenses/LICENSE-2.0
|||
||| Unless required by applicable law or agreed to in writing, software
||| distributed under the License is distributed on an "AS IS" BASIS,
||| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
||| See the License for the specific language governing permissions and
||| limitations under the License.
|||
||| This module defines a generalized algebraic data type, `MessageDescriptor`,
||| which is the strongly typed equivalent of the `Descriptor` protocol buffer.
||| Because types are first class values in Idris, we can construct a function
||| `interpMessage` that maps a `MessageDescriptor` to the *type* of structure
||| that the message describes.  The ability to treat types as first class
||| values, and to create generic functions allows us to do in Idris code much
||| of what is done with generated code for protocol buffer implementations in
||| other languages.
|||
||| We try to follow the naming conventions for protocol buffer descriptors
||| found in [descriptor.proto](https://developers.google.com/protocol-buffers/docs/reference/cpp/google.protobuf.descriptor.pb) with the exception that we use the name
||| `MessageDescriptor` instead of `Descriptor`.

module Protobuf.Core

import public Data.Fin
import public Data.Vect

%default total
%access public export

||| A `Label` is an annotation on a field that describes what kind of data type
||| the field represents.
data Label : Type where
  ||| `Optional` fields are represented as a `Maybe` type.
  Optional : Label
  ||| `Required` fields are represented as a single value.
  Required : Label
  ||| `Repeated` fields are represented as a `List`.
  Repeated : Label

||| A descriptor for a single value in an enum.
record EnumValueDescriptor where
  constructor MkEnumValueDescriptor
  name : String
  number : Int

||| A descriptor for an enum.
|||
||| k is the size of the enum
||| values is the possible values of the enum.  Note that enum values have a
||| numeric value that is possibly different from the index of the value in
||| `values`, e.g. `values` could be `[MkEnumValueDescriptor "five" 5]`.
record EnumDescriptor where
  constructor MkEnumDescriptor
  name : String
  values : Vect k EnumValueDescriptor

mutual
  ||| A descriptor that describes the type of a protocol buffer message.  This
  ||| describes the type of a protocol buffer message, which is a list of
  ||| fields.
  record MessageDescriptor where
    constructor MkMessageDescriptor
    name : String
    fields : Vect k FieldDescriptor

  ||| A descriptor that describes the type of a protocol buffer field. E.g.
  ||| `MkFieldDescriptor Optional PBDouble "temperature"` describes an
  ||| optional field whose name is "temperature" of type PBDouble, which is
  ||| interpreted Idris as a `Maybe double`.
  |||
  ||| label determines whether the field is optional, required or repeated.
  ||| ty is the type of a single value in the field.
  ||| name is the field's name, used in the text format.
  record FieldDescriptor where
    constructor MkFieldDescriptor
    label : Label
    ty : FieldValueDescriptor
    name : String
    number : Int

  ||| This does not correspond to anything in `descriptor.proto`, but is a
  ||| necessary definition here.  It describes the type of the value of a
  ||| protocol buffer required field, i.e. a single value.  This may be a
  ||| message, an enum or a primitive value.
  data FieldValueDescriptor : Type where
    PBDouble : FieldValueDescriptor
    PBFloat : FieldValueDescriptor
    PBInt32 : FieldValueDescriptor
    PBInt64 : FieldValueDescriptor
    PBUInt32 : FieldValueDescriptor
    PBUInt64 : FieldValueDescriptor
    PBSInt32 : FieldValueDescriptor
    PBSInt64 : FieldValueDescriptor
    PBFixed32 : FieldValueDescriptor
    PBFixed64 : FieldValueDescriptor
    PBSFixed32 : FieldValueDescriptor
    PBSFixed64 : FieldValueDescriptor
    PBBool : FieldValueDescriptor
    PBString : FieldValueDescriptor
    PBBytes : FieldValueDescriptor
    ||| A value which is itself a message.
    PBMessage : MessageDescriptor -> FieldValueDescriptor
    PBEnum : EnumDescriptor -> FieldValueDescriptor

interpEnum : EnumDescriptor -> Type
interpEnum (MkEnumDescriptor {k=k} _ _) = Fin k

-- NOTE: we implement Interp* as inductive types rather than as functions,
-- because these allow us to implement type classes such as Eq and Show.
--
-- E.g. if we had a function interpMessage : MessageDescriptor -> Type
-- it would not be possible to define a typeclass implementation
--
-- implementation Eq (interpMessage d) where ...
--
-- but it is possible to define
--
-- implementation Eq (InterpMessage d) where ...
--
-- Type classes also operate differently with respect to unfolding, though
-- that didn't pose any problems while working with functions.

mutual
  ||| Takes a message descriptor and returns the type that that message
  ||| descriptor describes.  This is implemented with the algebraic data type
  ||| `InterpMessage`.
  data InterpMessage : MessageDescriptor -> Type where
    MkMessage : InterpFields fields -> InterpMessage (MkMessageDescriptor name fields)

  ||| `InterpMessage` the algebraic data type used to implement `interpMessage`.
  ||| `InterpMessage fs` is a heterogeneous vector where where the type of each
  ||| element is `interpField` applied to the corresponding element in `fs`.
  ||| While it is possible to construct this type using `Data.HVect`, I haven't
  ||| found a way to satisfy the totality checker when using this approach.
  data InterpFields : (Vect k FieldDescriptor) -> Type where
    Nil  : InterpFields Nil
    (::) : (interpField f) -> (InterpFields m) -> InterpFields (f :: m)

  ||| Takes a field descriptor and returns the type that that field descriptor
  ||| describes.  This is a single value, a `Maybe` or a `List` depending on
  ||| the field's label.
  interpField : FieldDescriptor -> Type
  interpField (MkFieldDescriptor Optional value _ _) = Maybe (interpFieldValue value)
  interpField (MkFieldDescriptor Required value _ _) = interpFieldValue value
  interpField (MkFieldDescriptor Repeated value _ _) = List (interpFieldValue value)

  ||| Interprets a single field value, i.e. the value described by a required
  ||| field.  For primitive types this is a corresponding Idris primitive type.
  ||| For messages this is `interpMessage` for that message.  For enum's it is
  ||| `Fin k` where `k` is the number of values in the enum.
  interpFieldValue : FieldValueDescriptor -> Type
  interpFieldValue PBDouble = Double
  interpFieldValue PBFloat = Double
  interpFieldValue PBInt32 = Integer
  interpFieldValue PBInt64 = Integer
  interpFieldValue PBUInt32 = Integer
  interpFieldValue PBUInt64 = Integer
  interpFieldValue PBSInt32 = Integer
  interpFieldValue PBSInt64 = Integer
  interpFieldValue PBFixed32 = Integer
  interpFieldValue PBFixed64 = Integer
  interpFieldValue PBSFixed32 = Integer
  interpFieldValue PBSFixed64 = Integer
  interpFieldValue PBBool = Bool
  interpFieldValue PBString = String
  interpFieldValue PBBytes = String
  interpFieldValue (PBMessage m) = InterpMessage m
  interpFieldValue (PBEnum e) = interpEnum e


module Erlang.Conversion

import Erlang.Types

%default total

-- UNSAFE CONVERSIONS

-- IsErlType a -> IsErlType b

||| Cast a value from one type to another. (Caution: Wrong usage of this
||| function may result in run-time errors.)
|||
||| Both the target type and the type of the given value must be a member of `IsErlType`.
|||
||| `erlUnsafeCast` should only be used when you are sure the casting is safe.
||| In the common case, you should probably use the functions from `Erlang.Decode`.
export %inline
erlUnsafeCast : IsErlType from => (0 to : Type) -> IsErlType to => from -> to
erlUnsafeCast _ = believe_me

-- CONVERSION TO `ErlTerm`

||| Convert a value to an untyped Erlang value (`ErlTerm`).
public export
interface ToErlTerm a where
  toErlTerm : a -> ErlTerm

-- Values that conform to `IsErlType` are already proper Erlang values.
export %inline
IsErlType a => ToErlTerm a where
  toErlTerm = believe_me

export
ToErlTerm Bool where
  toErlTerm True = toErlTerm (MkAtom "true")
  toErlTerm False = toErlTerm (MkAtom "false")

-- CASTS

-- Bool -> ErlAtom

export
Cast Bool ErlAtom where
  cast True = MkAtom "true"
  cast False = MkAtom "false"

-- ErlNil/ErlCons -> List

export
Cast ErlNil (List a) where
  cast [] = []

export
Cast b (List a) => Cast (ErlCons a b) (List a) where
  cast (x :: y) = x :: cast y

-- ErlNil/ErlCons -> ErlList

export
Cast ErlNil (ErlList []) where
  cast [] = []

export
Cast b (ErlList bs) => Cast (ErlCons a b) (ErlList (a :: bs)) where
  cast (x :: y) = x :: cast y

-- ErlList -> List

export
Cast (ErlList []) (List a) where
  cast [] = []

export
Cast (ErlList as) (List a) => Cast (ErlList (a :: as)) (List a) where
  cast (x :: xs) = x :: cast xs

-- ErlList -> ErlNil/ErlCons

export
Cast (ErlList []) ErlNil where
  cast [] = []

export
Cast (ErlList [a]) (ErlCons a ErlNil) where
  cast (x :: xs) = x :: cast xs

export
Cast (ErlList as) (ErlCons b c) => Cast (ErlList (a :: as)) (ErlCons a (ErlCons b c)) where
  cast (x :: xs) = x :: cast xs

-- List -> ErlNil/ErlCons

||| Helper functions for `listToErlCons`.
export
0 ListToErlCons : List a -> Type
ListToErlCons [] = ErlNil
ListToErlCons (_ :: xs) = ErlCons a (ListToErlCons xs)

||| Convert `List` to `ErlNil` and `ErlCons`.
export
listToErlCons : (xs : List a) -> ListToErlCons xs
listToErlCons [] = []
listToErlCons (x :: xs) = x :: listToErlCons xs

-- List -> ErlList

||| Helper function for `listToErlList`.
export
0 ListToTypes : List a -> List Type
ListToTypes [] = []
ListToTypes (_ :: xs) = a :: ListToTypes xs

||| Convert `List` to `ErlList`.
export
listToErlList : (xs : List a) -> ErlList (ListToTypes xs)
listToErlList [] = []
listToErlList (x :: xs) = x :: listToErlList xs


module Lambdapants.Command.Parser

import Lambdapants.Command
import Lambdapants.Term
import Lambdapants.Term.Parser
import Lambdapants.Term.Reduction
import Lightyear
import Lightyear.Char
import Lightyear.Strings

symbolArg : Parser String
symbolArg = symbol

termArg : Parser Term
termArg = (char '(' *> appl <* char ')') <|>| map Var symbol

natArg : Parser Nat
natArg = (cast . pack) <$> some (satisfy isDigit)

stratArg : Parser Strategy
stratArg = norm <|> appl where
  norm : Parser Strategy
  norm = string "normal" *> pure Normal
  appl : Parser Strategy
  appl = string "applicative" *> pure Applicative

data ArgT
  = Arg0 Command
  | Arg1 (a -> Command)           (Parser a)
  | Arg2 (a -> b -> Command)      (Parser a) (Parser b)
  | Arg3 (a -> b -> c -> Command) (Parser a) (Parser b) (Parser c)

arity : ArgT -> Nat
arity (Arg0 _      ) = 0
arity (Arg1 _ _    ) = 1
arity (Arg2 _ _ _  ) = 2
arity (Arg3 _ _ _ _) = 3

commands : List (String, (ArgT, String))
commands =
  [ ("q"       , (Arg0 Quit                      , ""))
  , ("quit"    , (Arg0 Quit                      , ""))
  , ("limit"   , (Arg1 Limit natArg              , "<number>"))
  , ("u"       , (Arg1 Unset symbolArg           , "<symbol>"))
  , ("unset"   , (Arg1 Unset symbolArg           , "<symbol>"))
  , ("s"       , (Arg2 Set symbolArg termArg     , "<symbol> <term>"))
  , ("set"     , (Arg2 Set symbolArg termArg     , "<symbol> <term>"))
  , ("w"       , (Arg1 Whatis termArg            , "<term>"))
  , ("whatis"  , (Arg1 Whatis termArg            , "<term>"))
  , ("r"       , (Arg1 Reduce termArg            , "<term>"))
  , ("reduce"  , (Arg1 Reduce termArg            , "<term>"))
  , ("eq"      , (Arg2 Eq termArg termArg        , "<term> <term>"))
  , ("aq"      , (Arg2 AlphaEq termArg termArg   , "<term> <term>"))
  , ("env"     , (Arg0 Env                       , ""))
  , ("eval"    , (Arg1 Eval (opt stratArg)       , "[normal | applicative]"))
  , ("?"       , (Arg0 Help                      , ""))
  , ("h"       , (Arg0 Help                      , ""))
  , ("help"    , (Arg0 Help                      , "")) ]

args : ArgT -> Parser Command
args (Arg0 constr) = do
  spaces
  eof
  pure constr
args (Arg1 constr arg1) = do
  a <- arg1
  spaces
  eof
  pure (constr a)
args (Arg2 constr arg1 arg2) = do
  a <- arg1
  some space
  b <- arg2
  spaces
  eof
  pure (constr a b)
args (Arg3 constr arg1 arg2 arg3) = do
  a <- arg1
  some space
  b <- arg2
  some space
  c <- arg3
  spaces
  eof
  pure (constr a b c)

parseArgs : String -> String -> Either String Command
parseArgs command argstr =
  maybe (Left err) (uncurry compile) (lookup command commands)
where
  compile : ArgT -> String -> Either String Command
  compile argt hint = do
    case parse (args argt) argstr of
         Left _  => Left (if 0 == arity argt
                         then command ++ " is like Chuck Norris. \
                           \It takes no arguments."
                         else "Usage is :" ++ command ++ " " ++ hint)
         right => right
  err : String
  err = "\ESC[0;91mUnrecognized command: " ++ command ++ "\ESC[0m"

export
parseCmd : String -> Either String Command
parseCmd str = do
  case words str of
       (w :: ws) =>
            let args = unwords ws in
                if "!" == w
                   then Right (Shell args)
                   else parseArgs w args
       otherwise => Left "error"


module TyTTP.Core.Request

public export
record Request m u h a where
  constructor MkRequest
  method : m
  url : u
  headers : h
  body : a

export
Functor (Request m u h) where
  map f req = { body $= f } req