Fix bug #654

Auxiliary functions introduced in elaboration (e.g., through case splits and with clauses) now
have the same totality annotation as the function they're defined in.

Moved auxiliary function `findSetTotal` into `Context.idr` since it's
now used by `ProcessDef.idr` too.

Added a totality requirement argument to `checkClause` so that the
with-clause case could propagate it to the functions it generates in
elaboration.

Sandwhich the rhs elaboration in pattern matches with code that sets
the global, default, totality requirement to the current one, and
restores the previous default afterwards. It's a bit of a hacky way to
do it, but I don't think we have a better alternative with the current
design.

src/Core/Context.idr

@@ -288,6 +288,14 @@ export
 refersToRuntime : GlobalDef -> NameMap Bool
 refersToRuntime def = maybe empty id (refersToRuntimeM def)
 
+export
+findSetTotal : List DefFlag -> Maybe TotalReq
+findSetTotal [] = Nothing
+findSetTotal (SetTotal t :: _) = Just t
+findSetTotal (_ :: xs) = findSetTotal xs
+
+
+
 -- Label for array references
 export
 data Arr : Type where

src/TTImp/ProcessDecls.idr

@@ -95,11 +95,6 @@ checkTotalityOK n
                          _ => pure Nothing)
                    (pure . Just) err
 
-    findSetTotal : List DefFlag -> Maybe TotalReq
-    findSetTotal [] = Nothing
-    findSetTotal (SetTotal t :: _) = Just t
-    findSetTotal (_ :: xs) = findSetTotal xs
-
 -- Check totality of all the names added in the file, and return a list of
 -- totality errors.
 -- Do this at the end of processing a file (or a batch of definitions) since

src/TTImp/ProcessDef.idr

@@ -30,6 +30,7 @@ import Data.Either
 import Data.List
 import Data.NameMap
 import Data.Strings
+import Data.Maybe
 
 import Text.PrettyPrint.Prettyprinter
 
@@ -359,10 +360,11 @@ checkClause : {vars : _} ->
               {auto c : Ref Ctxt Defs} ->
               {auto m : Ref MD Metadata} ->
               {auto u : Ref UST UState} ->
-              (mult : RigCount) -> (vis : Visibility) -> (hashit : Bool) ->
+              (mult : RigCount) -> (vis : Visibility) ->
+              (totreq : TotalReq) -> (hashit : Bool) ->
               Int -> List ElabOpt -> NestedNames vars -> Env Term vars ->
               ImpClause -> Core (Either RawImp Clause)
-checkClause mult vis hashit n opts nest env (ImpossibleClause fc lhs)
+checkClause mult vis totreq hashit n opts nest env (ImpossibleClause fc lhs)
     = do lhs_raw <- lhsInCurrentNS nest lhs
          handleUnify
            (do autoimp <- isUnboundImplicits
@@ -387,7 +389,7 @@ checkClause mult vis hashit n opts nest env (ImpossibleClause fc lhs)
                            if !(impossibleErrOK defs err)
                               then pure (Left lhs_raw)
                               else throw (ValidCase fc env (Right err)))
-checkClause {vars} mult vis hashit n opts nest env (PatClause fc lhs_in rhs)
+checkClause {vars} mult vis totreq hashit n opts nest env (PatClause fc lhs_in rhs)
     = do (_, (vars'  ** (sub', env', nest', lhstm', lhsty'))) <-
              checkLHS False mult hashit n opts nest env fc lhs_in
          let rhsMode = if isErased mult then InType else InExpr
@@ -412,7 +414,7 @@ checkClause {vars} mult vis hashit n opts nest env (PatClause fc lhs_in rhs)
 
          pure (Right (MkClause env' lhstm' rhstm))
 -- TODO: (to decide) With is complicated. Move this into its own module?
-checkClause {vars} mult vis hashit n opts nest env (WithClause fc lhs_in wval_raw flags cs)
+checkClause {vars} mult vis totreq hashit n opts nest env (WithClause fc lhs_in wval_raw flags cs)
     = do (lhs, (vars'  ** (sub', env', nest', lhspat, reqty))) <-
              checkLHS False mult hashit n opts nest env fc lhs_in
          let wmode
@@ -474,8 +476,9 @@ checkClause {vars} mult vis hashit n opts nest env (WithClause fc lhs_in wval_ra
          log "declare.def.clause" 5 $ "Argument names " ++ show wargNames
 
          wname <- genWithName !(prettyName !(toFullNames (Resolved n)))
-         widx <- addDef wname (newDef fc wname (if isErased mult then erased else top)
+         widx <- addDef wname (record {flags $= (SetTotal totreq ::)}
-                                      vars wtype vis None)
+                                    (newDef fc wname (if isErased mult then erased else top)
+                                      vars wtype vis None))
          let rhs_in = apply (IVar fc wname)
                         (map (IVar fc) envns ++
                          map (maybe wval_raw (\pn => IVar fc (snd pn))) wargNames)
@@ -603,7 +606,7 @@ mkRunTime fc n
            let PMDef r cargs tree_ct _ pats = definition gdef
                 | _ => pure () -- not a function definition
            let ty = type gdef
-           -- Prepare RHS of definitions, by erasing 0-multiplicities, and
+           -- Prepare RHS of definitions, by erasing 0m-ultiplicities, and
            -- finding any applications to specialise (partially evaluate)
            pats' <- traverse (toErased (location gdef) (getSpec (flags gdef)))
                              pats
@@ -708,8 +711,17 @@ processDef opts nest env fc n_in cs_in
                        then erased
                        else linear
          nidx <- resolveName n
-         cs <- traverse (checkClause mult (visibility gdef)
+         
+         -- Set the default totality option to the current one, saving the global one
+         defaultTotality <- getDefaultTotalityOption
+         let treq = fromMaybe defaultTotality (findSetTotal (flags gdef))
+         setDefaultTotalityOption treq
+         
+         cs <- traverse (checkClause mult (visibility gdef) treq
                                      hashit nidx opts nest env) cs_in
+         -- Restore global default
+         setDefaultTotalityOption defaultTotality
+                                                                 
          let pats = map toPats (rights cs)
 
          (cargs ** (tree_ct, unreachable)) <-

tests/Main.idr

@@ -110,7 +110,7 @@ idrisTests
        "reg029", "reg030", "reg031", "reg032", "reg033", "reg034",
        -- Totality checking
        "total001", "total002", "total003", "total004", "total005",
-       "total006", "total007", "total008", "total009",
+       "total006", "total007", "total008", "total009", "total010",
        -- The 'with' rule
        "with001", "with002",
        -- with-disambiguation

tests/idris2/total010/PartialWith.idr

@@ -0,0 +1,10 @@
+%default total
+
+partial 
+foo : Nat -> Nat
+foo Z = Z
+
+partial
+bar : Nat -> Nat
+bar n with (0)
+ bar n | k = foo n

tests/idris2/total010/expected

@@ -0,0 +1 @@
+1/1: Building PartialWith (PartialWith.idr)

tests/idris2/total010/run

@@ -0,0 +1,3 @@
+$1 --no-color --console-width 0 PartialWith.idr --check
+
+rm -rf build