From 21d5ef7d91c8c12c4d46d5b6a6a7daa6d00843de Mon Sep 17 00:00:00 2001
From: Steve Dunham <dunhamsteve@gmail.com>
Date: Thu, 25 Jul 2024 04:51:57 -0700
Subject: [PATCH 1/4] [ new ] totality checking can look under constructors
 (#3328)

* [ total ] Consider (x :: zs) to be smaller than (x :: y :: zs)

* Expand RHS metas in totality checking
---
 CHANGELOG_NEXT.md                         |   3 +
 src/Core/Termination/CallGraph.idr        | 102 ++++++++++++++++------
 tests/idris2/total/total004/expected      |   2 +-
 tests/idris2/total/total025/Issue3317.idr |   5 ++
 tests/idris2/total/total025/Totality.idr  |  38 ++++++++
 tests/idris2/total/total025/expected      |   2 +
 tests/idris2/total/total025/run           |   4 +
 tests/ttimp/total002/expected             |   2 +-
 8 files changed, 128 insertions(+), 30 deletions(-)
 create mode 100644 tests/idris2/total/total025/Issue3317.idr
 create mode 100644 tests/idris2/total/total025/Totality.idr
 create mode 100644 tests/idris2/total/total025/expected
 create mode 100755 tests/idris2/total/total025/run

diff --git a/CHANGELOG_NEXT.md b/CHANGELOG_NEXT.md
index 283dbc7b50..66f2007fee 100644
--- a/CHANGELOG_NEXT.md
+++ b/CHANGELOG_NEXT.md
@@ -74,6 +74,9 @@ This CHANGELOG describes the merged but unreleased changes. Please see [CHANGELO
 * The compiler now parses `~x.fun` as unquoting `x` rather than `x.fun`
   and `~(f 5).fun` as unquoting `(f 5)` rather than `(f 5).fun`.
 
+* Totality checking will now look under data constructors, so `Just xs` will
+  be considered smaller than `Just (x :: xs)`.
+
 * LHS of `with`-applications are parsed as `PWithApp` instead of `PApp`. As a
   consequence, `IWithApp` appears in `TTImp` values in elaborator scripts instead
   of `IApp`, as it should have been.
diff --git a/src/Core/Termination/CallGraph.idr b/src/Core/Termination/CallGraph.idr
index 9b318ad3d1..2d809aedcc 100644
--- a/src/Core/Termination/CallGraph.idr
+++ b/src/Core/Termination/CallGraph.idr
@@ -1,11 +1,13 @@
 module Core.Termination.CallGraph
 
+import Core.Case.CaseTree
 import Core.Context
 import Core.Context.Log
 import Core.Env
 import Core.Normalise
 import Core.Value
 
+import Libraries.Data.IntMap
 import Libraries.Data.SparseMatrix
 
 import Data.String
@@ -134,56 +136,100 @@ mutual
                     else pure $ Ref fc Func n
         conIfGuarded tm = pure tm
 
-  knownOr : SizeChange -> Lazy SizeChange -> SizeChange
-  knownOr Unknown y = y
-  knownOr x _ = x
+  knownOr : Core SizeChange -> Core SizeChange -> Core SizeChange
+  knownOr x y = case !x of Unknown => y; _ => x
 
-  plusLazy : SizeChange -> Lazy SizeChange -> SizeChange
-  plusLazy Smaller _ = Smaller
-  plusLazy x y = x |+| y
+  plusLazy : Core SizeChange -> Core SizeChange -> Core SizeChange
+  plusLazy x y = case !x of Smaller => pure Smaller; x => pure $ x |+| !y
 
   -- Return whether first argument is structurally smaller than the second.
-  sizeCompare : Term vars -> -- RHS: term we're checking
+  sizeCompare : {auto defs : Defs} ->
+                Term vars -> -- RHS: term we're checking
                 Term vars -> -- LHS: argument it might be smaller than
-                SizeChange
+                Core SizeChange
 
-  sizeCompareCon : Term vars -> Term vars -> Bool
-  sizeCompareConArgs : Term vars -> List (Term vars) -> Bool
-  sizeCompareApp : Term vars -> Term vars -> SizeChange
+  sizeCompareCon : {auto defs : Defs} -> Term vars -> Term vars -> Core Bool
+  sizeCompareTyCon : {auto defs : Defs} -> Term vars -> Term vars -> Bool
+  sizeCompareConArgs : {auto defs : Defs} -> Term vars -> List (Term vars) -> Core Bool
+  sizeCompareApp : {auto defs : Defs} -> Term vars -> Term vars -> Core SizeChange
 
   sizeCompare s (Erased _ (Dotted t)) = sizeCompare s t
-  sizeCompare _ (Erased _ _) = Unknown -- incomparable!
+  sizeCompare _ (Erased _ _) = pure Unknown -- incomparable!
   -- for an as pattern, it's smaller if it's smaller than either part
   sizeCompare s (As _ _ p t)
       = knownOr (sizeCompare s p) (sizeCompare s t)
   sizeCompare (As _ _ p s) t
       = knownOr (sizeCompare p t) (sizeCompare s t)
+  -- if they're both metas, let sizeEq check if they're the same
+  sizeCompare s@(Meta _ _ _ _) t@(Meta _ _ _ _) = pure (if sizeEq s t then Same else Unknown)
+  -- otherwise try to expand RHS meta
+  sizeCompare s@(Meta n _ i args) t = do
+    Just gdef <- lookupCtxtExact (Resolved i) (gamma defs) | _ => pure Unknown
+    let (PMDef _ [] (STerm _ tm) _ _) = definition gdef | _ => pure Unknown
+    tm <- substMeta (embed tm) args zero []
+    sizeCompare tm t
+    where
+      substMeta : {0 drop, vs : _} ->
+                  Term (drop ++ vs) -> List (Term vs) ->
+                  SizeOf drop -> SubstEnv drop vs ->
+                  Core (Term vs)
+      substMeta (Bind bfc n (Lam _ c e ty) sc) (a :: as) drop env
+          = substMeta sc as (suc drop) (a :: env)
+      substMeta (Bind bfc n (Let _ c val ty) sc) as drop env
+          = substMeta (subst val sc) as drop env
+      substMeta rhs [] drop env = pure (substs drop env rhs)
+      substMeta rhs _ _ _ = throw (InternalError ("Badly formed metavar solution \{show n}"))
+
   sizeCompare s t
-     = if sizeCompareCon s t
-          then Smaller
-          else knownOr (sizeCompareApp s t) (if sizeEq s t then Same else Unknown)
+     = if sizeCompareTyCon s t then pure Same
+       else if !(sizeCompareCon s t)
+          then pure Smaller
+          else knownOr (sizeCompareApp s t) (pure $ if sizeEq s t then Same else Unknown)
+
+  -- consider two types the same size
+  sizeCompareTyCon s t =
+    let (f, args) = getFnArgs t in
+    let (g, args') = getFnArgs s in
+    case f of
+      Ref _ (TyCon _ _) cn => case g of
+        Ref _ (TyCon _ _) cn' => cn == cn'
+        _ => False
+      _ => False
 
   sizeCompareCon s t
       = let (f, args) = getFnArgs t in
         case f of
-             Ref _ (DataCon t a) cn => sizeCompareConArgs s args
-             _ => False
-
-  sizeCompareConArgs s [] = False
+             Ref _ (DataCon t a) cn =>
+                -- if s is smaller or equal to an arg, then it is smaller than t
+                if !(sizeCompareConArgs s args) then pure True
+                else let (g, args') = getFnArgs s in
+                    case g of
+                        Ref _ (DataCon t' a') cn' =>
+                                -- if s is a matching DataCon, applied to same number of args,
+                                -- no Unknown args, and at least one Smaller
+                                if cn == cn' && length args == length args'
+                                  then do
+                                       sizes <- traverse (uncurry sizeCompare) (zip args' args)
+                                       pure $ Smaller == foldl (|*|) Same sizes
+                                  else pure False
+                        _ => pure $ False
+             _ => pure False
+
+  sizeCompareConArgs s [] = pure False
   sizeCompareConArgs s (t :: ts)
-      = case sizeCompare s t of
+      = case !(sizeCompare s t) of
           Unknown => sizeCompareConArgs s ts
-          _ => True
+          _ => pure True
 
   sizeCompareApp (App _ f _) t = sizeCompare f t
-  sizeCompareApp _ t = Unknown
+  sizeCompareApp _ t = pure Unknown
 
-  sizeCompareAsserted : Maybe (Term vars) -> Term vars -> SizeChange
+  sizeCompareAsserted : {auto defs : Defs} -> Maybe (Term vars) -> Term vars -> Core SizeChange
   sizeCompareAsserted (Just s) t
-      = case sizeCompare s t of
+      = pure $ case !(sizeCompare s t) of
           Unknown => Unknown
           _ => Smaller
-  sizeCompareAsserted Nothing _ = Unknown
+  sizeCompareAsserted Nothing _ = pure Unknown
 
   -- if the argument is an 'assert_smaller', return the thing it's smaller than
   asserted : Name -> Term vars -> Maybe (Term vars)
@@ -200,9 +246,9 @@ mutual
   mkChange : Defs -> Name ->
              (pats : List (Term vars)) ->
              (arg : Term vars) ->
-             List SizeChange
+             Core (List SizeChange)
   mkChange defs aSmaller pats arg
-    = map (\p => plusLazy (sizeCompareAsserted (asserted aSmaller arg) p) (sizeCompare arg p)) pats
+    = traverse (\p => plusLazy (sizeCompareAsserted (asserted aSmaller arg) p) (sizeCompare arg p)) pats
 
   -- Given a name of a case function, and a list of the arguments being
   -- passed to it, update the pattern list so that it's referring to the LHS
@@ -310,7 +356,7 @@ mutual
               (do scs <- traverse (findSC defs env g pats) args
                   pure ([MkSCCall fn
                            (fromListList
-                                (map (mkChange defs aSmaller pats) args))
+                                !(traverse (mkChange defs aSmaller pats) args))
                            fc]
                            ++ concat scs))
 
diff --git a/tests/idris2/total/total004/expected b/tests/idris2/total/total004/expected
index 614ebee778..3a16dc8a3c 100644
--- a/tests/idris2/total/total004/expected
+++ b/tests/idris2/total/total004/expected
@@ -6,7 +6,7 @@ Main> Main.bar is total
 Main> Main.swapR is total
 Main> Main.loopy is possibly not terminating due to recursive path Main.loopy
 Main> Main.foom is total
-Main> Main.pfoom is possibly not terminating due to recursive path Main.pfoom
+Main> Main.pfoom is total
 Main> Main.even is total
 Main> Main.vtrans is possibly not terminating due to recursive path Main.vtrans -> Main.vtrans
 Main> Main.GTree is total
diff --git a/tests/idris2/total/total025/Issue3317.idr b/tests/idris2/total/total025/Issue3317.idr
new file mode 100644
index 0000000000..f2eaba1f3a
--- /dev/null
+++ b/tests/idris2/total/total025/Issue3317.idr
@@ -0,0 +1,5 @@
+total
+f : (a, List a) -> ()
+f (_, []) = ()
+f (x, _::xs) = f (x, xs)
+
diff --git a/tests/idris2/total/total025/Totality.idr b/tests/idris2/total/total025/Totality.idr
new file mode 100644
index 0000000000..a3830b3495
--- /dev/null
+++ b/tests/idris2/total/total025/Totality.idr
@@ -0,0 +1,38 @@
+%default total
+
+-- This one needed withHoles on tcOnly (instead of withArgHoles)
+-- or the solved implicit for the (ty : Type) on Maybe would not be
+-- substituted.
+length : Maybe (List a) -> Nat
+length Nothing = 0
+length (Just []) = 0
+length (Just (x :: xs)) = 1 + length (Just xs)
+
+one : List a -> Nat
+one (x :: y :: zs) = one (x :: zs)
+one _ = 0
+
+two : List a -> Nat
+two (a :: b :: c :: d :: es) = two (a :: c :: es)
+two _ = 0
+
+three : List a -> Nat
+three (a :: b :: c :: d :: es) = three (b :: c :: es)
+three _ = 0
+
+failing "not total"
+  four : List Nat -> Nat
+  four (a :: b :: c :: es) = four (b :: c :: a :: es)
+  four _ = 0
+
+-- also needs withHoles
+five : (List a, List a) -> List a
+five (a :: as, bs) = a :: five (as, bs)
+five ([], _) = []
+
+-- This is total, but not supported
+failing "not total"
+  six : (List a, List a) -> List a
+  six (a :: as, bs) = a :: six (bs, as)
+  six ([], _) = []
+
diff --git a/tests/idris2/total/total025/expected b/tests/idris2/total/total025/expected
new file mode 100644
index 0000000000..dd89c83aae
--- /dev/null
+++ b/tests/idris2/total/total025/expected
@@ -0,0 +1,2 @@
+1/1: Building Issue3317 (Issue3317.idr)
+1/1: Building Totality (Totality.idr)
diff --git a/tests/idris2/total/total025/run b/tests/idris2/total/total025/run
new file mode 100755
index 0000000000..369c11b9cb
--- /dev/null
+++ b/tests/idris2/total/total025/run
@@ -0,0 +1,4 @@
+. ../../../testutils.sh
+
+idris2 --check Issue3317.idr
+idris2 --check Totality.idr
diff --git a/tests/ttimp/total002/expected b/tests/ttimp/total002/expected
index 592c5ee5ac..eab9fde62c 100644
--- a/tests/ttimp/total002/expected
+++ b/tests/ttimp/total002/expected
@@ -5,5 +5,5 @@ Yaffle> Main.ack is total
 Yaffle> Main.foo is total
 Yaffle> Main.foo' is total
 Yaffle> Main.foom is total
-Yaffle> Main.pfoom is possibly not terminating due to recursive path Main.pfoom
+Yaffle> Main.pfoom is total
 Yaffle> Bye for now!

From 572cb981c77b509902e588f6c7b89425d849c54f Mon Sep 17 00:00:00 2001
From: Steve Dunham <dunhamsteve@gmail.com>
Date: Sat, 27 Jul 2024 14:08:26 -0700
Subject: [PATCH 2/4] [ fix ] Add fuel to totality checking under constructors

---
 docs/source/reference/pragmas.rst         |  7 +++
 src/Core/Context.idr                      |  5 ++
 src/Core/Options.idr                      |  3 +-
 src/Core/Termination/CallGraph.idr        | 75 +++++++++++++----------
 src/Idris/Desugar.idr                     |  1 +
 src/Idris/Parser.idr                      |  4 ++
 src/Idris/Syntax.idr                      |  1 +
 tests/idris2/total/total025/Issue3353.idr | 12 ++++
 tests/idris2/total/total025/expected      |  1 +
 tests/idris2/total/total025/run           |  1 +
 10 files changed, 76 insertions(+), 34 deletions(-)
 create mode 100644 tests/idris2/total/total025/Issue3353.idr

diff --git a/docs/source/reference/pragmas.rst b/docs/source/reference/pragmas.rst
index 76612e2edb..2d7f58dd1d 100644
--- a/docs/source/reference/pragmas.rst
+++ b/docs/source/reference/pragmas.rst
@@ -62,6 +62,13 @@ definition.
 Set how deep nested ambiguous names can be before Idris gives up. The default is 3, increashing this
 will effect compiler performance. For more details, see :ref:`ambiguous-name-resolution`.
 
+``%totality_depth``
+-------------------
+
+Set the number of matching constructors Idris will look under when checking totality.  For instance
+`Just xs` is smaller than `Just (x :: xs)` if Idris looks under the matching constructor. The default
+value is 5.  Increasing the value may slow down totality checking.
+
 ``%auto_implicit_depth``
 ------------------------
 
diff --git a/src/Core/Context.idr b/src/Core/Context.idr
index ce7256c79e..0655399fcf 100644
--- a/src/Core/Context.idr
+++ b/src/Core/Context.idr
@@ -2252,6 +2252,11 @@ setAmbigLimit : {auto c : Ref Ctxt Defs} ->
                 Nat -> Core ()
 setAmbigLimit max = update Ctxt { options->elabDirectives->ambigLimit := max }
 
+export
+setTotalLimit : {auto c : Ref Ctxt Defs} ->
+                Nat -> Core ()
+setTotalLimit max = update Ctxt { options->elabDirectives->totalLimit := max }
+
 export
 setAutoImplicitLimit : {auto c : Ref Ctxt Defs} ->
                        Nat -> Core ()
diff --git a/src/Core/Options.idr b/src/Core/Options.idr
index 57d1773049..45470f3fb5 100644
--- a/src/Core/Options.idr
+++ b/src/Core/Options.idr
@@ -127,6 +127,7 @@ record ElabDirectives where
   ambigLimit : Nat
   autoImplicitLimit : Nat
   nfThreshold : Nat
+  totalLimit : Nat
   --
   -- produce traditional (prefix) record projections,
   -- in addition to postfix (dot) projections
@@ -232,7 +233,7 @@ defaultSession = MkSessionOpts False CoveringOnly False False Chez [] 1000 False
 
 export
 defaultElab : ElabDirectives
-defaultElab = MkElabDirectives True True CoveringOnly 3 50 25 True
+defaultElab = MkElabDirectives True True CoveringOnly 3 50 25 5 True
 
 -- FIXME: This turns out not to be reliably portable, since different systems
 -- may have tools with the same name but different required arugments. We
diff --git a/src/Core/Termination/CallGraph.idr b/src/Core/Termination/CallGraph.idr
index 2d809aedcc..c9e442ee2a 100644
--- a/src/Core/Termination/CallGraph.idr
+++ b/src/Core/Termination/CallGraph.idr
@@ -5,6 +5,7 @@ import Core.Context
 import Core.Context.Log
 import Core.Env
 import Core.Normalise
+import Core.Options
 import Core.Value
 
 import Libraries.Data.IntMap
@@ -144,30 +145,31 @@ mutual
 
   -- Return whether first argument is structurally smaller than the second.
   sizeCompare : {auto defs : Defs} ->
+                Nat -> -- backtracking fuel
                 Term vars -> -- RHS: term we're checking
                 Term vars -> -- LHS: argument it might be smaller than
                 Core SizeChange
 
-  sizeCompareCon : {auto defs : Defs} -> Term vars -> Term vars -> Core Bool
+  sizeCompareCon : {auto defs : Defs} -> Nat -> Term vars -> Term vars -> Core Bool
   sizeCompareTyCon : {auto defs : Defs} -> Term vars -> Term vars -> Bool
-  sizeCompareConArgs : {auto defs : Defs} -> Term vars -> List (Term vars) -> Core Bool
-  sizeCompareApp : {auto defs : Defs} -> Term vars -> Term vars -> Core SizeChange
+  sizeCompareConArgs : {auto defs : Defs} -> Nat -> Term vars -> List (Term vars) -> Core Bool
+  sizeCompareApp : {auto defs : Defs} -> Nat -> Term vars -> Term vars -> Core SizeChange
 
-  sizeCompare s (Erased _ (Dotted t)) = sizeCompare s t
-  sizeCompare _ (Erased _ _) = pure Unknown -- incomparable!
+  sizeCompare fuel s (Erased _ (Dotted t)) = sizeCompare fuel s t
+  sizeCompare fuel _ (Erased _ _) = pure Unknown -- incomparable!
   -- for an as pattern, it's smaller if it's smaller than either part
-  sizeCompare s (As _ _ p t)
-      = knownOr (sizeCompare s p) (sizeCompare s t)
-  sizeCompare (As _ _ p s) t
-      = knownOr (sizeCompare p t) (sizeCompare s t)
+  sizeCompare fuel s (As _ _ p t)
+      = knownOr (sizeCompare fuel s p) (sizeCompare fuel s t)
+  sizeCompare fuel (As _ _ p s) t
+      = knownOr (sizeCompare fuel p t) (sizeCompare fuel s t)
   -- if they're both metas, let sizeEq check if they're the same
-  sizeCompare s@(Meta _ _ _ _) t@(Meta _ _ _ _) = pure (if sizeEq s t then Same else Unknown)
+  sizeCompare fuel s@(Meta _ _ _ _) t@(Meta _ _ _ _) = pure (if sizeEq s t then Same else Unknown)
   -- otherwise try to expand RHS meta
-  sizeCompare s@(Meta n _ i args) t = do
+  sizeCompare fuel s@(Meta n _ i args) t = do
     Just gdef <- lookupCtxtExact (Resolved i) (gamma defs) | _ => pure Unknown
     let (PMDef _ [] (STerm _ tm) _ _) = definition gdef | _ => pure Unknown
     tm <- substMeta (embed tm) args zero []
-    sizeCompare tm t
+    sizeCompare fuel tm t
     where
       substMeta : {0 drop, vs : _} ->
                   Term (drop ++ vs) -> List (Term vs) ->
@@ -180,11 +182,11 @@ mutual
       substMeta rhs [] drop env = pure (substs drop env rhs)
       substMeta rhs _ _ _ = throw (InternalError ("Badly formed metavar solution \{show n}"))
 
-  sizeCompare s t
+  sizeCompare fuel s t
      = if sizeCompareTyCon s t then pure Same
-       else if !(sizeCompareCon s t)
+       else if !(sizeCompareCon fuel s t)
           then pure Smaller
-          else knownOr (sizeCompareApp s t) (pure $ if sizeEq s t then Same else Unknown)
+          else knownOr (sizeCompareApp fuel s t) (pure $ if sizeEq s t then Same else Unknown)
 
   -- consider two types the same size
   sizeCompareTyCon s t =
@@ -196,40 +198,46 @@ mutual
         _ => False
       _ => False
 
-  sizeCompareCon s t
+  sizeCompareProdConArgs : {auto defs : Defs} -> Nat -> List (Term vars) -> List (Term vars) -> Core SizeChange
+  sizeCompareProdConArgs _ [] [] = pure Same
+  sizeCompareProdConArgs fuel (x :: xs) (y :: ys) =
+    case !(sizeCompare fuel x y) of
+      Unknown => pure Unknown
+      t => (t |*|) <$> sizeCompareProdConArgs fuel xs ys
+  sizeCompareProdConArgs _ _ _ = pure Unknown
+
+  sizeCompareCon fuel s t
       = let (f, args) = getFnArgs t in
         case f of
              Ref _ (DataCon t a) cn =>
                 -- if s is smaller or equal to an arg, then it is smaller than t
-                if !(sizeCompareConArgs s args) then pure True
+                if !(sizeCompareConArgs (minus fuel 1) s args) then pure True
                 else let (g, args') = getFnArgs s in
-                    case g of
-                        Ref _ (DataCon t' a') cn' =>
+                    case (fuel, g) of
+                        (S k, Ref _ (DataCon t' a') cn') => do
                                 -- if s is a matching DataCon, applied to same number of args,
                                 -- no Unknown args, and at least one Smaller
                                 if cn == cn' && length args == length args'
-                                  then do
-                                       sizes <- traverse (uncurry sizeCompare) (zip args' args)
-                                       pure $ Smaller == foldl (|*|) Same sizes
+                                  then (Smaller ==) <$> sizeCompareProdConArgs k args' args
                                   else pure False
                         _ => pure $ False
              _ => pure False
 
-  sizeCompareConArgs s [] = pure False
-  sizeCompareConArgs s (t :: ts)
-      = case !(sizeCompare s t) of
-          Unknown => sizeCompareConArgs s ts
+  sizeCompareConArgs _ s [] = pure False
+  sizeCompareConArgs fuel s (t :: ts)
+      = case !(sizeCompare fuel s t) of
+          Unknown => sizeCompareConArgs fuel s ts
           _ => pure True
 
-  sizeCompareApp (App _ f _) t = sizeCompare f t
-  sizeCompareApp _ t = pure Unknown
+  sizeCompareApp fuel (App _ f _) t = sizeCompare fuel f t
+  sizeCompareApp _ _ t = pure Unknown
 
-  sizeCompareAsserted : {auto defs : Defs} -> Maybe (Term vars) -> Term vars -> Core SizeChange
-  sizeCompareAsserted (Just s) t
-      = pure $ case !(sizeCompare s t) of
+  sizeCompareAsserted : {auto defs : Defs} -> Nat -> Maybe (Term vars) -> Term vars -> Core SizeChange
+  sizeCompareAsserted fuel (Just s) t
+      = pure $ case !(sizeCompare fuel s t) of
           Unknown => Unknown
           _ => Smaller
-  sizeCompareAsserted Nothing _ = pure Unknown
+  sizeCompareAsserted _ Nothing _ = pure Unknown
 
   -- if the argument is an 'assert_smaller', return the thing it's smaller than
   asserted : Name -> Term vars -> Maybe (Term vars)
@@ -248,7 +256,8 @@ mutual
              (arg : Term vars) ->
              Core (List SizeChange)
   mkChange defs aSmaller pats arg
-    = traverse (\p => plusLazy (sizeCompareAsserted (asserted aSmaller arg) p) (sizeCompare arg p)) pats
+    = let fuel = defs.options.elabDirectives.totalLimit
+      in traverse (\p => plusLazy (sizeCompareAsserted fuel (asserted aSmaller arg) p) (sizeCompare fuel arg p)) pats
 
   -- Given a name of a case function, and a list of the arguments being
   -- passed to it, update the pattern list so that it's referring to the LHS
diff --git a/src/Idris/Desugar.idr b/src/Idris/Desugar.idr
index 02b41723e1..749d0cbdc2 100644
--- a/src/Idris/Desugar.idr
+++ b/src/Idris/Desugar.idr
@@ -1319,6 +1319,7 @@ mutual
              PrefixRecordProjections b => do
                pure [IPragma fc [] (\nest, env => setPrefixRecordProjections b)]
              AmbigDepth n => pure [IPragma fc [] (\nest, env => setAmbigLimit n)]
+             TotalityDepth n => pure [IPragma fc [] (\next, env => setTotalLimit n)]
              AutoImplicitDepth n => pure [IPragma fc [] (\nest, env => setAutoImplicitLimit n)]
              NFMetavarThreshold n => pure [IPragma fc [] (\nest, env => setNFThreshold n)]
              SearchTimeout n => pure [IPragma fc [] (\nest, env => setSearchTimeout n)]
diff --git a/src/Idris/Parser.idr b/src/Idris/Parser.idr
index fd941c347a..6410401be4 100644
--- a/src/Idris/Parser.idr
+++ b/src/Idris/Parser.idr
@@ -1466,6 +1466,10 @@ directive fname indents
          b <- onoff
          atEnd indents
          pure (PrefixRecordProjections b)
+  <|> do decoratedPragma fname "totality_depth"
+         lvl <- decorate fname Keyword $ intLit
+         atEnd indents
+         pure (TotalityDepth (fromInteger lvl))
   <|> do decoratedPragma fname "ambiguity_depth"
          lvl <- decorate fname Keyword $ intLit
          atEnd indents
diff --git a/src/Idris/Syntax.idr b/src/Idris/Syntax.idr
index 3973b1cb4e..595ccfc6e1 100644
--- a/src/Idris/Syntax.idr
+++ b/src/Idris/Syntax.idr
@@ -362,6 +362,7 @@ mutual
        LazyOn : Bool -> Directive
        UnboundImplicits : Bool -> Directive
        AmbigDepth : Nat -> Directive
+       TotalityDepth: Nat -> Directive
        PairNames : Name -> Name -> Name -> Directive
        RewriteName : Name -> Name -> Directive
        PrimInteger : Name -> Directive
diff --git a/tests/idris2/total/total025/Issue3353.idr b/tests/idris2/total/total025/Issue3353.idr
new file mode 100644
index 0000000000..db26dc2380
--- /dev/null
+++ b/tests/idris2/total/total025/Issue3353.idr
@@ -0,0 +1,12 @@
+import Data.Fin
+
+covering
+g : Fin 64 -> Unit
+g FZ      = ()
+g (FS i') = g $ weaken i'
+
+total
+g' : Fin 64 -> Unit
+g' FZ        = ()
+g' i@(FS i') = g' $ assert_smaller i $ weaken i'
+
diff --git a/tests/idris2/total/total025/expected b/tests/idris2/total/total025/expected
index dd89c83aae..5b972ab256 100644
--- a/tests/idris2/total/total025/expected
+++ b/tests/idris2/total/total025/expected
@@ -1,2 +1,3 @@
 1/1: Building Issue3317 (Issue3317.idr)
 1/1: Building Totality (Totality.idr)
+1/1: Building Issue3353 (Issue3353.idr)
diff --git a/tests/idris2/total/total025/run b/tests/idris2/total/total025/run
index 369c11b9cb..5858578fd2 100755
--- a/tests/idris2/total/total025/run
+++ b/tests/idris2/total/total025/run
@@ -2,3 +2,4 @@
 
 idris2 --check Issue3317.idr
 idris2 --check Totality.idr
+idris2 --check Issue3353.idr

From 1caa7cff4f73aaba2c5706be57c3d83e81d26024 Mon Sep 17 00:00:00 2001
From: Steve Dunham <dunhamsteve@gmail.com>
Date: Sat, 27 Jul 2024 14:55:08 -0700
Subject: [PATCH 3/4] [ test ] addition tests for totality checking

---
 tests/idris2/total/total025/Pragma.idr   | 14 ++++++++++++++
 tests/idris2/total/total025/Totality.idr | 12 ++++++++++++
 tests/idris2/total/total025/expected     |  1 +
 tests/idris2/total/total025/run          |  1 +
 4 files changed, 28 insertions(+)
 create mode 100644 tests/idris2/total/total025/Pragma.idr

diff --git a/tests/idris2/total/total025/Pragma.idr b/tests/idris2/total/total025/Pragma.idr
new file mode 100644
index 0000000000..383cee8c06
--- /dev/null
+++ b/tests/idris2/total/total025/Pragma.idr
@@ -0,0 +1,14 @@
+%default total
+
+%totality_depth 1
+
+failing "not total"
+  one : List a -> Nat
+  one (a :: b :: c :: es) = one (a :: b :: es)
+  one _ = 0
+
+%totality_depth 3
+
+two : List a -> Nat
+two (a :: b :: c :: es) = two (a :: b :: es)
+two _ = 0
diff --git a/tests/idris2/total/total025/Totality.idr b/tests/idris2/total/total025/Totality.idr
index a3830b3495..9f6427a6c3 100644
--- a/tests/idris2/total/total025/Totality.idr
+++ b/tests/idris2/total/total025/Totality.idr
@@ -36,3 +36,15 @@ failing "not total"
   six (a :: as, bs) = a :: six (bs, as)
   six ([], _) = []
 
+
+failing "not total"
+  -- If we didn't check all of the arguments of MkTuple for
+  -- Same/Smaller, then this would be incorrectly accepted as total
+  first : (List Nat, List Nat) -> Nat
+  second : (List Nat, List Nat) -> Nat
+
+  first (x :: xs, ys) = second (xs, Z :: ys)
+  first _ = Z
+
+  second (xs, y :: ys) = first (1 :: xs, ys)
+  second _ = Z
diff --git a/tests/idris2/total/total025/expected b/tests/idris2/total/total025/expected
index 5b972ab256..27e347939f 100644
--- a/tests/idris2/total/total025/expected
+++ b/tests/idris2/total/total025/expected
@@ -1,3 +1,4 @@
 1/1: Building Issue3317 (Issue3317.idr)
 1/1: Building Totality (Totality.idr)
 1/1: Building Issue3353 (Issue3353.idr)
+1/1: Building Pragma (Pragma.idr)
diff --git a/tests/idris2/total/total025/run b/tests/idris2/total/total025/run
index 5858578fd2..e80910c86e 100755
--- a/tests/idris2/total/total025/run
+++ b/tests/idris2/total/total025/run
@@ -3,3 +3,4 @@
 idris2 --check Issue3317.idr
 idris2 --check Totality.idr
 idris2 --check Issue3353.idr
+idris2 --check Pragma.idr

From 4c5e079237523001575cba7e5546b5cb59753adb Mon Sep 17 00:00:00 2001
From: Steve Dunham <dunhamsteve@gmail.com>
Date: Wed, 9 Oct 2024 21:05:54 -0700
Subject: [PATCH 4/4] [ fix ] compare arguments on type constructors

---
 src/Core/Termination/CallGraph.idr | 25 +++++++++++++------------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/src/Core/Termination/CallGraph.idr b/src/Core/Termination/CallGraph.idr
index c9e442ee2a..a29f716768 100644
--- a/src/Core/Termination/CallGraph.idr
+++ b/src/Core/Termination/CallGraph.idr
@@ -151,7 +151,7 @@ mutual
                 Core SizeChange
 
   sizeCompareCon : {auto defs : Defs} -> Nat -> Term vars -> Term vars -> Core Bool
-  sizeCompareTyCon : {auto defs : Defs} -> Term vars -> Term vars -> Bool
+  sizeCompareTyCon : {auto defs : Defs} -> Nat -> Term vars -> Term vars -> Core Bool
   sizeCompareConArgs : {auto defs : Defs} -> Nat -> Term vars -> List (Term vars) -> Core Bool
   sizeCompareApp : {auto defs : Defs} -> Nat -> Term vars -> Term vars -> Core SizeChange
 
@@ -183,21 +183,11 @@ mutual
       substMeta rhs _ _ _ = throw (InternalError ("Badly formed metavar solution \{show n}"))
 
   sizeCompare fuel s t
-     = if sizeCompareTyCon s t then pure Same
+     = if !(sizeCompareTyCon fuel s t) then pure Same
        else if !(sizeCompareCon fuel s t)
           then pure Smaller
           else knownOr (sizeCompareApp fuel s t) (pure $ if sizeEq s t then Same else Unknown)
 
-  -- consider two types the same size
-  sizeCompareTyCon s t =
-    let (f, args) = getFnArgs t in
-    let (g, args') = getFnArgs s in
-    case f of
-      Ref _ (TyCon _ _) cn => case g of
-        Ref _ (TyCon _ _) cn' => cn == cn'
-        _ => False
-      _ => False
-
   sizeCompareProdConArgs : {auto defs : Defs} -> Nat -> List (Term vars) -> List (Term vars) -> Core SizeChange
   sizeCompareProdConArgs _ [] [] = pure Same
   sizeCompareProdConArgs fuel (x :: xs) (y :: ys) =
@@ -206,6 +196,17 @@ mutual
       t => (t |*|) <$> sizeCompareProdConArgs fuel xs ys
   sizeCompareProdConArgs _ _ _ = pure Unknown
 
+  sizeCompareTyCon fuel s t =
+    let (f, args) = getFnArgs t in
+    let (g, args') = getFnArgs s in
+    case f of
+      Ref _ (TyCon _ _) cn => case g of
+        Ref _ (TyCon _ _) cn' => if cn == cn'
+            then (Unknown /=) <$> sizeCompareProdConArgs fuel args' args
+            else pure False
+        _ => pure False
+      _ => pure False
+
   sizeCompareCon fuel s t
       = let (f, args) = getFnArgs t in
         case f of