Refactor %builtin (#1803)

* Stub for future 'identity' optimisation I plan to add this later, but I'm using for now for NaturalToInteger and IntegerToNatural * Refactor `%builtin` fixes #1799 - automatically optimise all Natural shaped things - NaturalToInteger and IntegerToNatural now use new `Identity` flag (internal use only for now) which signals the function is identity at runtime * Use NaturalToInteger and IntegerToNatural for Nat and Fin Also define show fin in terms of finToInteger, for speed * Fix name handling for %builtin * [ tests ] fixes + #1799 * remove %builtin from libs Add back after next version * Use resolved names where convenient
2024-12-21 10:41:59 +03:00 · 2021-08-03 14:19:17 +01:00 · 2021-08-03 14:19:17 +01:00 · 1669fc351b
commit 1669fc351b
parent ed0c8d60bc
21 changed files with 275 additions and 265 deletions
--- a/idris2api.ipkg
+++ b/idris2api.ipkg
@ -12,6 +12,7 @@ modules =
    Compiler.Common,
    Compiler.CompileExpr,
    Compiler.Generated,
+    Compiler.Identity,
    Compiler.Inline,
    Compiler.LambdaLift,
    Compiler.Separate,
--- a/libs/base/Data/Fin.idr
+++ b/libs/base/Data/Fin.idr
@ -60,11 +60,6 @@ finToNat : Fin n -> Nat
 finToNat FZ = Z
 finToNat (FS k) = S $ finToNat k

-
-export
-Show (Fin n) where
-  show = show . finToNat
-
 ||| `finToNat` is injective
 export
 finToNatInjective : (fm : Fin k) -> (fn : Fin k) -> (finToNat fm) = (finToNat fn) -> fm = fn
@ -84,6 +79,12 @@ finToInteger : Fin n -> Integer
 finToInteger FZ     = 0
 finToInteger (FS k) = 1 + finToInteger k

+-- %builtin NaturalToInteger Data.Fin.finToInteger
+
+export
+Show (Fin n) where
+  show = show . finToInteger
+
 export
 Cast (Fin n) Integer where
    cast = finToInteger
@ -168,6 +169,8 @@ fromInteger : (x : Integer) -> {n : Nat} ->
 fromInteger {n} x {prf} with (integerToFin x n)
  fromInteger {n} x {prf = ItIsJust} | Just y = y

+-- %builtin IntegerToNatural Data.Fin.fromInteger
+
 ||| Convert an Integer to a Fin in the required bounds/
 ||| This is essentially a composition of `mod` and `fromInteger`
 public export
--- a/libs/prelude/Prelude/Types.idr
+++ b/libs/prelude/Prelude/Types.idr
@ -39,6 +39,8 @@ integerToNat x
       then Z
       else S (assert_total (integerToNat (prim__sub_Integer x 1)))

+-- %builtin IntegerToNatural Prelude.Types.integerToNat
+
 -- Define separately so we can spot the name when optimising Nats
 ||| Add two natural numbers.
 ||| @ x the number to case-split on
@ -99,6 +101,8 @@ natToInteger (S k) = 1 + natToInteger k
                         -- integer (+) may be non-linear in second
                         -- argument

+-- %builtin NaturalToInteger Prelude.Types.natToInteger
+
 ||| Counts the number of elements that satify a predicate.
 public export
 count : (Foldable t) => (predicate : a -> Bool) -> (t a) -> Nat
--- a/src/Compiler/CompileExpr.idr
+++ b/src/Compiler/CompileExpr.idr
@ -211,93 +211,51 @@ natHack =

 -- get all transformation from %builtin pragmas
 builtinMagic : Ref Ctxt Defs => Core (forall vars. CExp vars -> CExp vars)
-builtinMagic = do
-    defs <- get Ctxt
-    let b = defs.builtinTransforms
-    let nats = foldMap builtinMagicNat $ values b.natTyNames
-    let natToInts = map natToIntMagic $ toList b.natToIntegerFns
-    let intToNats = map intToNatMagic $ toList b.integerToNatFns
-    pure $ magic $ natHack ++ nats ++ natToInts ++ intToNats
-  where
-    builtinMagicNat : NatBuiltin -> List Magic
-    builtinMagicNat cons =
-        [ MagicCCon cons.zero 0
-             (\ fc, [] => CPrimVal fc (BI 0))
-        , MagicCCon cons.succ 1
-             (\ fc, [k] => CApp fc (CRef fc (UN "prim__add_Integer")) [CPrimVal fc (BI 1), k])
-        ] -- TODO: add builtin pragmas for Nat related functions (to/from Integer, add, mult, minus, compare)
-    natToIntMagic : (Name, NatToInt) -> Magic
-    natToIntMagic (fn, MkNatToInt arity natIdx) =
-        MagicCRef fn arity
-            (\ _, _, args => index natIdx args)
-    intToNatMagic : (Name, IntToNat) -> Magic
-    intToNatMagic (fn, MkIntToNat arity intIdx) =
-        MagicCRef fn arity
-            (\ fc, fc', args => CApp fc (CRef fc' (NS typesNS (UN "prim__integerToNat"))) [index intIdx args])
+builtinMagic = pure $ magic natHack

-isNatCon : (zeroMap : NameMap ZERO) ->
-           (succMap : NameMap SUCC) ->
-           Name -> Bool
-isNatCon zs ss n = isJust (lookup n zs) || isJust (lookup n ss)
+data NextSucc : Type where
+newSuccName : {auto s : Ref NextSucc Int} -> Core Name
+newSuccName = do
+    x <- get NextSucc
+    put NextSucc $ x + 1
+    pure $ MN "succ" x

-natBranch : (zeroMap : NameMap ZERO) ->
-           (succMap : NameMap SUCC) ->
-           CConAlt vars -> Bool
-natBranch zs ss (MkConAlt n _ _ _ _) = isNatCon zs ss n
+natBranch :  CConAlt vars -> Bool
+natBranch (MkConAlt n ZERO _ _ _) = True
+natBranch (MkConAlt n SUCC _ _ _) = True
+natBranch _ = False

-trySBranch :
-    (succMap : NameMap SUCC) ->
-    CExp vars ->
-    CConAlt vars ->
-    Maybe (CExp vars)
-trySBranch ss n (MkConAlt nm _ _ [arg] sc)
-  = do guard $ isJust $ lookup nm ss
-       let fc = getFC n
-       pure (CLet fc arg True (magic__natUnsuc fc fc [n]) sc)
-trySBranch _ _ _ = Nothing
+trySBranch : CExp vars -> CConAlt vars -> Maybe (CExp vars)
+trySBranch n (MkConAlt nm SUCC _ [arg] sc)
+    = Just (CLet (getFC n) arg True (magic__natUnsuc (getFC n) (getFC n) [n]) sc)
+trySBranch _ _ = Nothing

-tryZBranch :
-    (zeroMap : NameMap ZERO) ->
-    CConAlt vars ->
-    Maybe (CExp vars)
-tryZBranch zs (MkConAlt n _ _ [] sc)
-   = do guard $ isJust $ lookup n zs
-        pure sc
-tryZBranch _ _ = Nothing
+tryZBranch : CConAlt vars -> Maybe (CExp vars)
+tryZBranch (MkConAlt n ZERO _ [] sc) = Just sc
+tryZBranch _ = Nothing

-getSBranch :
-    (succMap : NameMap SUCC) ->
-    CExp vars ->
-    List (CConAlt vars) ->
-    Maybe (CExp vars)
-getSBranch ss n [] = Nothing
-getSBranch ss n (x :: xs) = trySBranch ss n x <+> getSBranch ss n xs
+getSBranch : CExp vars -> List (CConAlt vars) -> Maybe (CExp vars)
+getSBranch n [] = Nothing
+getSBranch n (x :: xs) = trySBranch n x <+> getSBranch n xs

-getZBranch :
-    (zeroMap : NameMap ZERO) ->
-    List (CConAlt vars) ->
-    Maybe (CExp vars)
-getZBranch zs [] = Nothing
-getZBranch zs (x :: xs) = tryZBranch zs x <+> getZBranch zs xs
+getZBranch : List (CConAlt vars) -> Maybe (CExp vars)
+getZBranch [] = Nothing
+getZBranch (x :: xs) = tryZBranch x <+> getZBranch xs

 -- Rewrite case trees on Nat to be case trees on Integer
-builtinNatTree' : (zeroMap : NameMap ZERO) ->
-                  (succMap : NameMap SUCC) ->
-                  CExp vars -> CExp vars
-builtinNatTree' zs ss (CConCase fc sc alts def)
-   = if any (natBranch zs ss) alts
-        then let defb = maybe (CCrash fc "Nat case not covered") id def
-                 scase = maybe defb id (getSBranch ss sc alts)
-                 zcase = maybe defb id (getZBranch zs alts) in
-                 CConstCase fc sc [MkConstAlt (BI 0) zcase] (Just scase)
-        else CConCase fc sc alts def
-builtinNatTree' zs ss t = t
-
-builtinNatTree : Ref Ctxt Defs => Core (CExp vars -> CExp vars)
-builtinNatTree = do
-    defs <- get Ctxt
-    let b = defs.builtinTransforms
-    pure $ builtinNatTree' b.natZNames b.natSNames
+builtinNatTree : {auto s : Ref NextSucc Int} -> CExp vars -> Core (CExp vars)
+builtinNatTree (CConCase fc sc@(CLocal _ _) alts def)
+   = pure $ if any natBranch alts
+               then let defb = fromMaybe (CCrash fc "Nat case not covered") def
+                        salt = maybe defb id (getSBranch sc alts)
+                        zalt = maybe defb id (getZBranch alts) in
+                        CConstCase fc sc [MkConstAlt (BI 0) zalt] (Just salt)
+               else CConCase fc sc alts def
+builtinNatTree (CConCase fc sc alts def)
+    = do x <- newSuccName
+         pure $ CLet fc x True sc
+                !(builtinNatTree $ CConCase fc (CLocal fc First) (map weaken alts) (map weaken def))
+builtinNatTree t = pure t

 enumTree : CExp vars -> CExp vars
 enumTree (CConCase fc sc alts def)
@ -330,6 +288,7 @@ dconFlag n
 mutual
  toCExpTm : {vars : _} ->
             {auto c : Ref Ctxt Defs} ->
+             {auto s : Ref NextSucc Int} ->
             (magic : forall vars. CExp vars -> CExp vars) ->
             Name -> Term vars ->
             Core (CExp vars)
@ -341,6 +300,9 @@ mutual
           fl <- dconFlag cn
           case fl of
                ENUM => pure $ CPrimVal fc (I tag)
+                ZERO => pure $ CPrimVal fc (BI 0)
+                SUCC => do x <- newSuccName
+                           pure $ CLam fc x $ COp fc (Add IntegerType) [CPrimVal fc (BI 1), CLocal fc First]
                _ => pure $ CCon fc cn fl (Just tag) []
  toCExpTm m n (Ref fc (TyCon tag arity) fn)
      = pure $ CCon fc fn TYCON Nothing []
@ -382,6 +344,7 @@ mutual

  toCExp : {vars : _} ->
           {auto c : Ref Ctxt Defs} ->
+           {auto s : Ref NextSucc Int} ->
           (magic : forall vars. CExp vars -> CExp vars) ->
           Name -> Term vars ->
           Core (CExp vars)
@ -404,6 +367,7 @@ mutual
 mutual
  conCases : {vars : _} ->
             {auto c : Ref Ctxt Defs} ->
+             {auto s : Ref NextSucc Int} ->
             Name -> List (CaseAlt vars) ->
             Core (List (CConAlt vars))
  conCases n [] = pure []
@ -432,6 +396,7 @@ mutual

  constCases : {vars : _} ->
               {auto c : Ref Ctxt Defs} ->
+               {auto s : Ref NextSucc Int} ->
               Name -> List (CaseAlt vars) ->
               Core (List (CConstAlt vars))
  constCases n [] = pure []
@ -449,6 +414,7 @@ mutual
  -- once.
  getNewType : {vars : _} ->
               {auto c : Ref Ctxt Defs} ->
+               {auto s : Ref NextSucc Int} ->
               FC -> CExp vars ->
               Name -> List (CaseAlt vars) ->
               Core (Maybe (CExp vars))
@ -500,6 +466,7 @@ mutual

  getDef : {vars : _} ->
           {auto c : Ref Ctxt Defs} ->
+           {auto s : Ref NextSucc Int} ->
           Name -> List (CaseAlt vars) ->
           Core (Maybe (CExp vars))
  getDef n [] = pure Nothing
@ -511,6 +478,7 @@ mutual

  toCExpTree : {vars : _} ->
               {auto c : Ref Ctxt Defs} ->
+               {auto s : Ref NextSucc Int} ->
               Name -> CaseTree vars ->
               Core (CExp vars)
  toCExpTree n alts@(Case _ x scTy (DelayCase ty arg sc :: rest))
@ -524,6 +492,7 @@ mutual

  toCExpTree' : {vars : _} ->
                {auto c : Ref Ctxt Defs} ->
+                {auto s : Ref NextSucc Int} ->
                Name -> CaseTree vars ->
                Core (CExp vars)
  toCExpTree' n (Case _ x scTy alts@(ConCase _ _ _ _ :: _))
@ -535,8 +504,8 @@ mutual
               def <- getDef n alts
               if isNil cases
                  then pure (fromMaybe (CErased fc) def)
-                  else pure $ enumTree $ !builtinNatTree $
-                            CConCase fc (CLocal fc x) cases def
+                  else pure $ enumTree !(builtinNatTree $
+                            CConCase fc (CLocal fc x) cases def)
  toCExpTree' n (Case _ x scTy alts@(DelayCase _ _ _ :: _))
      = throw (InternalError "Unexpected DelayCase")
  toCExpTree' n (Case fc x scTy alts@(ConstCase _ _ :: _))
@ -731,6 +700,7 @@ toCDef n ty _ None
    = pure $ MkError $ CCrash emptyFC ("Encountered undefined name " ++ show !(getFullName n))
 toCDef n ty erased (PMDef pi args _ tree _)
    = do let (args' ** p) = mkSub 0 args erased
+         s <- newRef NextSucc 0
         comptree <- toCExpTree n tree
         pure $ toLam (externalDecl pi) $ if isNil erased
            then MkFun args comptree
@ -794,6 +764,7 @@ compileExp : {auto c : Ref Ctxt Defs} ->
             ClosedTerm -> Core (CExp [])
 compileExp tm
    = do m <- builtinMagic
+         s <- newRef NextSucc 0
         exp <- toCExp m (UN "main") tm
         pure exp

--- a/src/Compiler/Identity.idr
+++ b/src/Compiler/Identity.idr
@ -0,0 +1,36 @@
+module Compiler.Identity
+
+import Compiler.CompileExpr
+import Core.Context
+import Core.Context.Log
+
+getArg : FC -> Nat -> (args : List Name) -> Maybe (CExp args)
+getArg _ _ [] = Nothing
+getArg fc Z (a :: _) = Just $ CLocal fc First
+getArg fc (S k) (_ :: as) = weaken <$> getArg fc k as
+
+idCDef : Nat -> CDef -> Maybe CDef
+idCDef idx (MkFun args exp) = MkFun args <$> getArg (getFC exp) idx args
+idCDef _ def = Just def
+
+export
+setIdentity : Ref Ctxt Defs => Name -> Core ()
+setIdentity fn = do
+    log "compiler.identity" 10 $ "processing identity flag for: " ++ show fn
+    defs <- get Ctxt
+    Just gdef <- lookupCtxtExact fn defs.gamma
+        | Nothing => pure ()
+    log "compiler.identity" 25 $ "flags for " ++ show fn ++ ": " ++ show gdef.flags
+    let Just flg@(Identity idx) = find isId gdef.flags
+        | _ => pure ()
+    log "compiler.identity" 5 $ "found identity flag for: " ++ show fn ++ ", " ++ show idx
+                             ++ "\n\told def: " ++ show gdef.compexpr
+    let Just cdef = the _ $ gdef.compexpr >>= idCDef idx
+        | Nothing => pure ()
+    log "compiler.identity" 15 $ "new def: " ++ show cdef
+    unsetFlag EmptyFC fn flg -- other optimisations might mess with argument counts
+    setCompiled fn cdef
+  where
+    isId : DefFlag -> Bool
+    isId (Identity _) = True
+    isId _ = False
--- a/src/Compiler/Inline.idr
+++ b/src/Compiler/Inline.idr
@ -2,6 +2,7 @@ module Compiler.Inline

 import Compiler.CaseOpts
 import Compiler.CompileExpr
+import Compiler.Identity

 import Core.CompileExpr
 import Core.Context
@ -522,6 +523,7 @@ compileAndInlineAll
         cns <- filterM nonErased ns

         traverse_ compileDef cns
+         traverse_ setIdentity cns
         transform 3 cns -- number of rounds to run transformations.
                         -- This seems to be the point where not much useful
                         -- happens any more.
--- a/src/Core/Binary.idr
+++ b/src/Core/Binary.idr
@ -31,7 +31,7 @@ import public Libraries.Utils.Binary
 ||| (Increment this when changing anything in the data format)
 export
 ttcVersion : Int
-ttcVersion = 58
+ttcVersion = 59

 export
 checkTTCVersion : String -> Int -> Int -> Core ()
--- a/src/Core/CompileExpr.idr
+++ b/src/Core/CompileExpr.idr
@ -24,6 +24,8 @@ data ConInfo = DATACON -- normal data constructor
             | NOTHING -- nothing of an option shaped thing
             | JUST -- just of an option shaped thing
             | RECORD -- record constructor (no tag)
+             | ZERO
+             | SUCC

 export
 Show ConInfo where
@ -35,6 +37,8 @@ Show ConInfo where
  show NOTHING = "[nothing]"
  show JUST    = "[just]"
  show RECORD  = "[record]"
+  show ZERO    = "[zero]"
+  show SUCC    = "[succ]"

 export
 Eq ConInfo where
@ -46,6 +50,8 @@ Eq ConInfo where
  NOTHING == NOTHING = True
  JUST == JUST = True
  RECORD == RECORD = True
+  ZERO == ZERO = True
+  SUCC == SUCC = True
  _ == _ = False

 mutual
--- a/src/Core/Context.idr
+++ b/src/Core/Context.idr
@ -199,6 +199,9 @@ data DefFlag
    | ConType ConInfo
         -- Is it a special type of constructor, e.g. a nil or cons shaped
         -- thing, that can be compiled specially?
+    | Identity Nat
+         -- Is it the identity function at runtime?
+         -- The nat represents which argument the function evaluates to

 export
 Eq DefFlag where
@ -212,6 +215,7 @@ Eq DefFlag where
    (==) (PartialEval x) (PartialEval y) = x == y
    (==) AllGuarded AllGuarded = True
    (==) (ConType x) (ConType y) = x == y
+    (==) (Identity x) (Identity y) = x == y
    (==) _ _ = False

 export
@ -226,6 +230,7 @@ Show DefFlag where
  show (PartialEval _) = "partialeval"
  show AllGuarded = "allguarded"
  show (ConType ci) = "contype " ++ show ci
+  show (Identity x) = "identity " ++ show x

 public export
 data SizeChange = Smaller | Same | Unknown
@ -669,58 +674,6 @@ Show BuiltinType where
    show NaturalToInteger = "NaturalToInteger"
    show IntegerToNatural = "IntegerToNatural"

-- Token types to make it harder to get the constructor names
-- the wrong way round.
-public export data ZERO = MkZERO
-public export data SUCC = MkSUCC
-
-||| Record containing names of 'Nat'-like constructors.
-public export
-record NatBuiltin where
-    constructor MkNatBuiltin
-    zero : Name
-    succ : Name
-
-||| Record containing information about a NatToInteger function.
-public export
-record NatToInt where
-    constructor MkNatToInt
-    natToIntArity : Nat -- total number of arguments
-    natIdx : Fin natToIntArity -- index into arguments of the 'Nat'-like argument
-
-||| Record containing information about a IntegerToNat function.
-public export
-record IntToNat where
-    constructor MkIntToNat
-    intToNatArity : Nat
-    intIdx : Fin intToNatArity
-
-||| Rewrite rules for %builtin pragmas
-||| Seperate to 'Transform' because it must also modify case statements
-||| behaviour should remain the same after this transform
-public export
-record BuiltinTransforms where
-    constructor MkBuiltinTransforms
-    natTyNames : NameMap NatBuiltin -- map from Nat-like names to their constructors
-    natZNames : NameMap ZERO -- set of Z-like names
-    natSNames : NameMap SUCC -- set of S-like names
-    natToIntegerFns : NameMap NatToInt -- set of functions to transform to `id`
-    integerToNatFns : NameMap IntToNat -- set of functions to transform to `max 0`
-
-- TODO: After next release remove nat from here and use %builtin pragma instead
-initBuiltinTransforms : BuiltinTransforms
-initBuiltinTransforms =
-    let type = NS typesNS (UN "Nat")
-        zero = NS typesNS (UN "Z")
-        succ = NS typesNS (UN "S")
-    in MkBuiltinTransforms
-        { natTyNames = singleton type (MkNatBuiltin {zero, succ})
-        , natZNames = singleton zero MkZERO
-        , natSNames = singleton succ MkSUCC
-        , natToIntegerFns = empty
-        , integerToNatFns = empty
-        }
-
 export
 getFnName : Transform -> Maybe Name
 getFnName (MkTransform _ _ app _)
@ -1087,10 +1040,6 @@ record Defs where
     -- ^ A mapping from names to transformation rules which update applications
     -- of that name
  saveTransforms : List (Name, Transform)
-  builtinTransforms : BuiltinTransforms
-     -- ^ A mapping from names to transformations resulting from a %builtin pragma
-     -- seperate to `transforms` because these must always fire globally so run these
-     -- when compiling to `CExp`.
  namedirectives : NameMap (List String)
  ifaceHash : Int
  importHashes : List (Namespace, Int)
@ -1161,7 +1110,6 @@ initDefs
           , saveAutoHints = []
           , transforms = empty
           , saveTransforms = []
-           , builtinTransforms = initBuiltinTransforms
           , namedirectives = empty
           , ifaceHash = 5381
           , importHashes = []
--- a/src/Core/Hash.idr
+++ b/src/Core/Hash.idr
@ -402,7 +402,8 @@ Hashable ConInfo where
    NOTHING => h `hashWithSalt` 5
    JUST => h `hashWithSalt` 6
    RECORD => h `hashWithSalt` 7
-
+    ZERO => h `hashWithSalt` 8
+    SUCC => h `hashWithSalt` 9

 mutual
  export
--- a/src/Core/Options/Log.idr
+++ b/src/Core/Options/Log.idr
@ -53,6 +53,7 @@ knownTopics = [
    ("compile.casetree.measure", Just "Log the node counts of each runtime case tree."),
    ("compile.casetree.pick", Nothing),
    ("compile.casetree.partition", Nothing),
+    ("compiler.identity", Nothing),
    ("compiler.inline.eval", Nothing),
    ("compiler.interpreter", Nothing),
    ("compiler.refc", Nothing),
--- a/src/Core/TTC.idr
+++ b/src/Core/TTC.idr
@ -661,6 +661,8 @@ TTC ConInfo where
  toBuf b NOTHING = tag 5
  toBuf b JUST = tag 6
  toBuf b RECORD = tag 7
+  toBuf b ZERO = tag 8
+  toBuf b SUCC = tag 9

  fromBuf b
      = case !getTag of
@ -672,6 +674,8 @@ TTC ConInfo where
             5 => pure NOTHING
             6 => pure JUST
             7 => pure RECORD
+             8 => pure ZERO
+             9 => pure SUCC
             _ => corrupt "ConInfo"

 mutual
@ -1011,6 +1015,7 @@ TTC DefFlag where
  toBuf b (PartialEval x) = tag 9 -- names not useful any more
  toBuf b AllGuarded = tag 10
  toBuf b (ConType ci) = do tag 11; toBuf b ci
+  toBuf b (Identity x) = do tag 12; toBuf b x

  fromBuf b
      = case !getTag of
@ -1024,6 +1029,7 @@ TTC DefFlag where
             9 => pure (PartialEval [])
             10 => pure AllGuarded
             11 => do ci <- fromBuf b; pure (ConType ci)
+             12 => do x <- fromBuf b; pure (Identity x)
             _ => corrupt "DefFlag"

 export
--- a/src/TTImp/ProcessBuiltin.idr
+++ b/src/TTImp/ProcessBuiltin.idr
@ -11,6 +11,7 @@ import Core.CaseTree
 import Core.Core
 import Core.Context
 import Core.Context.Log
+import Core.CompileExpr
 import Core.Env
 import Core.Metadata
 import Core.TT
@ -78,33 +79,28 @@ getFirstNEType tm = case getNEArgs tm of
    arg :: _ => Just arg

 ||| Get the index of the first non-erased argument if it exists.
-getNEIndex : (arity : Nat) -> Term vars -> Maybe (Fin arity)
-getNEIndex ar (Bind _ x b tm) = case b of
-    Let _ _ val _ => getNEIndex ar $ subst {x} val tm
+getNEIndex : Term vars -> Maybe Nat
+getNEIndex (Bind _ x b tm) = case b of
+    Let _ _ val _ => getNEIndex $ subst {x} val tm
    Pi _ mul _ arg => if isErased mul
-        then getNEIndex ar tm >>= Libs.strengthen . FS
-        else natToFin 0 ar
+        then getNEIndex tm
+        else Just 0
    _ => Nothing
-getNEIndex _ _ = Nothing
+getNEIndex _ = Nothing

 ||| Get the index of all non-erased Integer argument.
-getNEIntegerIndex : (arity : Nat) -> Term vars -> Maybe (List (Fin arity))
-getNEIntegerIndex ar (Bind _ x b tm) = case b of
-    Let _ _ val _ => getNEIntegerIndex ar $ subst {x} val tm
-    Pi _ mul _ arg => if isRigOther mul && isInteger arg
-        then case natToFin 0 ar of
-            Nothing => Nothing
-            Just idx => map (Prelude.(::) idx) $ go ar tm
-        else go ar tm
+getNEIntegerIndex : Term vars -> Maybe (List Nat)
+getNEIntegerIndex (Bind _ x b tm) = case b of
+    Let _ _ val _ => getNEIntegerIndex $ subst {x} val tm
+    Pi _ mul _ arg => if not (isErased mul) && isInteger arg
+        then Prelude.(::) 0 . map (+ 1) <$> getNEIntegerIndex tm
+        else getNEIntegerIndex tm
    _ => Nothing
  where
    isInteger : forall vars. Term vars -> Bool
    isInteger (PrimVal _ IntegerType) = True
    isInteger _ = False
-    go : forall vars. (sarity : Nat) -> Term vars -> Maybe (List (Fin sarity))
-    go (S ar) tm = map FS <$> getNEIntegerIndex ar tm
-    go Z _ = Just []
-getNEIntegerIndex _ _ = Just []
+getNEIntegerIndex _ = Just []

 ||| Do the terms match ignoring arguments to type constructors.
 termConMatch : Term vs -> Term vs' -> Bool
@ -143,26 +139,43 @@ isStrict (PrimVal _ _) = True
 isStrict (Erased _ _) = True
 isStrict (TType _) = True

-getNatBuiltin : Ref Ctxt Defs => Name -> Core (Maybe NatBuiltin)
-getNatBuiltin n = do
-    n' <- getFullName n
-    lookup n' . natTyNames . builtinTransforms <$> get Ctxt
+||| Get the name and definition of a list of names.
+getConsGDef :
+    Ref Ctxt Defs =>
+    FC -> (cons : List Name) ->
+    Core $ List (Name, GlobalDef)
+getConsGDef fc cons = do
+    defs <- get Ctxt
+    let c = defs.gamma
+    for cons $ \n => do
+        [(n', _, gdef)] <- lookupCtxtName n c
+            | [] => undefinedName fc n
+            | ns => throw $ AmbiguousName fc $ (\(n, _, _) => n) <$> ns
+        pure (n', gdef)

-||| Get the name and arity (of non-erased arguments only) of a list of names.
-||| `cons` should all be data constructors (`DCon`) otherwise it will throw an error.
-getConsGDef : Context -> FC -> (cons : List Name) -> Core $ List (Name, GlobalDef)
-getConsGDef c fc = traverse \n => do
-    [(n', _, gdef)] <- lookupCtxtName n c
-        | [] => throw $ UndefinedName fc n
-        | ns => throw $ AmbiguousName fc $ (\(n, _, _) => n) <$> ns
-    pure (n', gdef)
+isNatural : Ref Ctxt Defs => FC ->Name -> Core Bool
+isNatural fc n = do
+    defs <- get Ctxt
+    Just gdef <- lookupCtxtExact n defs.gamma
+        | Nothing => undefinedName EmptyFC n
+    let TCon _ _ _ _ _ _ cons _ = gdef.definition
+        | _ => pure False
+    consDefs <- getConsGDef fc cons
+    pure $ all hasNatFlag consDefs
+  where
+    isNatFlag : DefFlag -> Bool
+    isNatFlag (ConType ZERO) = True
+    isNatFlag (ConType SUCC) = True
+    isNatFlag _ = False
+    hasNatFlag : (Name, GlobalDef) -> Bool
+    hasNatFlag (_, gdef) = any isNatFlag gdef.flags

 ||| Check a list of constructors has exactly
 ||| 1 'Z'-like constructor
 ||| and 1 `S`-like constructor, which has type `ty -> ty` or `ty arg -> `ty (f arg)`.
-checkNatCons : Context -> (cons : List (Name, GlobalDef)) -> (dataType : Name) -> FC -> Core NatBuiltin
+checkNatCons : Context -> (cons : List (Name, GlobalDef)) -> (dataType : Name) -> FC -> Core ()
 checkNatCons c cons ty fc = case !(foldr checkCon (pure (Nothing, Nothing)) cons) of
-    (Just zero, Just succ) => pure $ MkNatBuiltin {zero, succ}
+    (Just zero, Just succ) => pure ()
    (Nothing, _) => throw $ GenericMsg fc $ "No 'Z'-like constructors for " ++ show ty ++ "."
    (_, Nothing) => throw $ GenericMsg fc $ "No 'S'-like constructors for " ++ show ty ++ "."
  where
@ -204,115 +217,81 @@ checkNatCons c cons ty fc = case !(foldr checkCon (pure (Nothing, Nothing)) cons
                    pure (zero, Just n)
            _ => throw $ GenericMsg fc $ "Constructor " ++ show n ++ " doesn't match any pattern for Natural."

-addBuiltinNat :
-    Ref Ctxt Defs =>
-    (ty : Name) -> NatBuiltin -> Core ()
-addBuiltinNat type cons = do
-    log "builtin.Natural.addTransform" 10
-        $ "Add %builtin Natural transform for " ++ show type ++ "."
-    update Ctxt $ record
-        { builtinTransforms.natTyNames $= insert type cons
-        , builtinTransforms.natZNames $= insert cons.zero MkZERO
-        , builtinTransforms.natSNames $= insert cons.succ MkSUCC
-        }
-
-addNatToInteger :
-    Ref Ctxt Defs =>
-    (fn : Name) ->
-    NatToInt ->
-    Core ()
-addNatToInteger fn nToI = do
-    log "builtin.NaturalToInteger.addTransforms" 10
-        $ "Add %builtin NaturalToInteger transform for " ++ show fn ++ "."
-    update Ctxt $ record
-        { builtinTransforms.natToIntegerFns $= insert fn nToI
-        }
-
 ||| Check a `%builtin Natural` pragma is correct.
 processBuiltinNatural :
    Ref Ctxt Defs =>
-    Defs -> FC -> Name -> Core ()
-processBuiltinNatural ds fc name = do
+    FC -> Name -> Core ()
+processBuiltinNatural fc name = do
+    ds <- get Ctxt
    log "builtin.Natural" 5 $ "Processing %builtin Natural " ++ show name ++ "."
    [(n, _, gdef)] <- lookupCtxtName name ds.gamma
        | [] => undefinedName fc name
        | ns => throw $ AmbiguousName fc $ (\(n, _, _) => n) <$> ns
+    False <- isNatural fc n
+        | True => pure ()
    let TCon _ _ _ _ _ _ dcons _ = gdef.definition
        | def => throw $ GenericMsg fc
            $ "Expected a type constructor, found " ++ showDefType def ++ "."
-    cons <- getConsGDef ds.gamma fc dcons
-    cons <- checkNatCons ds.gamma cons n fc
-    zero <- getFullName cons.zero
-    succ <- getFullName cons.succ
-    addBuiltinNat n $ MkNatBuiltin {zero, succ}
+    cons <- getConsGDef fc dcons
+    checkNatCons ds.gamma cons n fc

 ||| Check a `%builtin NaturalToInteger` pragma is correct.
 processNatToInteger :
    Ref Ctxt Defs =>
-    Defs -> FC -> Name -> Core ()
-processNatToInteger ds fc fn = do
-    log "builtin.NaturalToInteger" 5 $ "Processing %builtin NaturalToInteger " ++ show fn ++ "."
-    [(n, _, gdef)] <- lookupCtxtName fn ds.gamma
+    FC -> Name -> Core ()
+processNatToInteger fc fn = do
+    let show_fn = show fn
+    ds <- get Ctxt
+    log "builtin.NaturalToInteger" 5 $ "Processing %builtin NaturalToInteger " ++ show_fn ++ "."
+    [(_ , i, gdef)] <- lookupCtxtName fn ds.gamma
        | [] => undefinedName fc fn
        | ns => throw $ AmbiguousName fc $ (\(n, _, _) => n) <$> ns
    let PMDef _ args _ cases _ = gdef.definition
        | def => throw $ GenericMsg fc
            $ "Expected function definition, found " ++ showDefType def ++ "."
    type <- toFullNames gdef.type
-    logTerm "builtin.NaturalToInteger" 25 ("Type of " ++ show fn) type
+    logTerm "builtin.NaturalToInteger" 25 ("Type of " ++ show_fn) type
    let [(_ ** arg)] = getNEArgs type
-        | [] => throw $ GenericMsg fc $ "No arguments found for " ++ show n ++ "."
-        | _ => throw $ GenericMsg fc $ "More than 1 non-erased arguments found for " ++ show n ++ "."
+        | [] => throw $ GenericMsg fc $ "No arguments found for " ++ show_fn ++ "."
+        | _ => throw $ GenericMsg fc $ "More than 1 non-erased arguments found for " ++ show_fn ++ "."
    let Just tyCon = getTypeCons arg
        | Nothing => throw $ GenericMsg fc
-            $ "No type constructor found for non-erased arguement of " ++ show n ++ "."
-    Just _ <- getNatBuiltin tyCon
-        | Nothing => throw $ GenericMsg fc $ "Non-erased argument is not a 'Nat'-like type."
-    let arity = length $ getTypeArgs type
-    let Just natIdx = getNEIndex arity type
+            $ "No type constructor found for non-erased arguement of " ++ show_fn ++ "."
+    True <- isNatural fc tyCon
+        | False => throw $ GenericMsg fc $ "Non-erased argument is not a 'Nat'-like type."
+    let Just natIdx = getNEIndex type
        | Nothing => throw $ InternalError "Couldn't find non-erased argument."
-    addNatToInteger n (MkNatToInt arity natIdx)
-
-addIntegerToNat :
-    Ref Ctxt Defs =>
-    (fn : Name) ->
-    IntToNat ->
-    Core ()
-addIntegerToNat fn iToN = do
-    log "builtin.IntegerToNatural.addTransforms" 10
-        $ "Add %builtin IntegerToNatural transform for " ++ show fn ++ "."
-    update Ctxt $ record
-        { builtinTransforms.integerToNatFns $= insert fn iToN
-        }
+    setFlag fc (Resolved i) $ Identity natIdx

 ||| Check a `%builtin IntegerToNatural` pragma is correct.
 processIntegerToNat :
    Ref Ctxt Defs =>
-    Defs -> FC -> Name -> Core ()
-processIntegerToNat ds fc fn = do
-    log "builtin.IntegerToNatural" 5 $ "Processing %builtin IntegerToNatural " ++ show fn ++ "."
-    [(n, _, gdef)] <- lookupCtxtName fn ds.gamma
+    FC -> Name -> Core ()
+processIntegerToNat fc fn = do
+    let show_fn = show fn
+    ds <- get Ctxt
+    log "builtin.IntegerToNatural" 5 $ "Processing %builtin IntegerToNatural " ++ show_fn ++ "."
+    [(_, i, gdef)] <- lookupCtxtName fn ds.gamma
        | [] => undefinedName fc fn
        | ns => throw $ AmbiguousName fc $ (\(n, _, _) => n) <$> ns
    type <- toFullNames gdef.type
-    let PMDef _ args _ cases _ = gdef.definition
+    let PMDef _ _ _ _ _ = gdef.definition
        | def => throw $ GenericMsg fc
            $ "Expected function definition, found " ++ showDefType def ++ "."
-    let arity = length $ getTypeArgs type
-    logTerm "builtin.IntegerToNatural" 25 ("Type of " ++ show fn) type
-    let Just [intIdx] = getNEIntegerIndex arity type
-        | Just [] => throw $ GenericMsg fc $ "No unrestricted arguments of type `Integer` found for " ++ show n ++ "."
-        | Just _ => throw $ GenericMsg fc $ "More than one unrestricted arguments of type `Integer` found for " ++ show n ++ "."
+    logTerm "builtin.IntegerToNatural" 25 ("Type of " ++ show_fn) type
+    let Just [intIdx] = getNEIntegerIndex type
+        | Just [] => throw $ GenericMsg fc $ "No unrestricted arguments of type `Integer` found for " ++ show_fn ++ "."
+        | Just _ => throw $ GenericMsg fc $ "More than one unrestricted arguments of type `Integer` found for " ++ show_fn ++ "."
        | Nothing => throw $ InternalError
-            $ "Unexpected arity while processing %builtin IntegerToNatural " ++ show n ++ " (getNEIntegerIndex returned Nothing)"
+            $ "Unexpected arity while processing %builtin IntegerToNatural " ++ show_fn ++ " (getNEIntegerIndex returned Nothing)"
    let Just (_ ** retTy) = getReturnType type
        | Nothing => throw $ InternalError $ "Unexpected type " ++ show type
    let Just retCon = getTypeCons retTy
        | Nothing => throw $ GenericMsg fc
-            $ "No type constructor found for return type of " ++ show n ++ "."
-    Just _ <- getNatBuiltin retCon
-        | Nothing => throw $ GenericMsg fc $ "Return type is not a 'Nat'-like type"
-    addIntegerToNat n (MkIntToNat arity intIdx)
+            $ "No type constructor found for return type of " ++ show_fn ++ "."
+    True <- isNatural fc retCon
+        | False => throw $ GenericMsg fc $ "Return type is not a 'Nat'-like type"
+    setFlag fc (Resolved i) $ Identity intIdx

 ||| Check a `%builtin` pragma is correct.
 export
@ -320,8 +299,7 @@ processBuiltin :
    Ref Ctxt Defs =>
    NestedNames vars -> Env Term vars -> FC -> BuiltinType -> Name -> Core ()
 processBuiltin nest env fc type name = do
-    ds <- get Ctxt
    case type of
-        BuiltinNatural => processBuiltinNatural ds fc name
-        NaturalToInteger => processNatToInteger ds fc name
-        IntegerToNatural => processIntegerToNat ds fc name
+        BuiltinNatural => processBuiltinNatural fc name
+        NaturalToInteger => processNatToInteger fc name
+        IntegerToNatural => processIntegerToNat fc name
--- a/src/TTImp/ProcessData.idr
+++ b/src/TTImp/ProcessData.idr
@ -19,6 +19,7 @@ import TTImp.Elab
 import TTImp.TTImp
 import TTImp.Utils

+import Data.DPair
 import Data.List
 import Libraries.Data.NameMap

@ -117,9 +118,6 @@ checkCon {vars} opts nest env vis tn_in tn (MkImpTy fc _ cn_in ty_raw)
              _ => pure ()
         pure (MkCon fc cn !(getArity defs [] fullty) fullty)

-conName : Constructor -> Name
-conName (MkCon _ cn _ _) = cn
-
 -- Get the indices of the constructor type (with non-constructor parts erased)
 getIndexPats : {auto c : Ref Ctxt Defs} ->
               ClosedTerm -> Core (List (NF []))
@ -254,6 +252,19 @@ hasArgs Z (Bind _ _ (Pi _ c _ _) sc)
         else False
 hasArgs Z _ = True

+-- get the first non-erased argument
+firstArg : Term vs -> Maybe (Exists Term)
+firstArg (Bind _ _ (Pi _ c _ val) sc)
+    = if isErased c
+         then firstArg sc
+         else Just $ Evidence _ val
+firstArg tm = Nothing
+
+typeCon : Term vs -> Maybe Name
+typeCon (Ref _ (TyCon _ _) n) = Just n
+typeCon (App _ fn _) = typeCon fn
+typeCon _ = Nothing
+
 shaped : {auto c : Ref Ctxt Defs} ->
         (forall vs . Term vs -> Bool) ->
         List Constructor -> Core (Maybe Name)
@ -322,10 +333,36 @@ calcRecord fc [c]
         pure True
 calcRecord _ _ = pure False

+-- has two constructors
+-- - ZERO: 0 args
+-- - SUCC: 1 arg, of same type
+calcNaty : {auto c : Ref Ctxt Defs} ->
+           FC -> Name -> List Constructor -> Core Bool
+calcNaty fc tyCon cs@[_, _]
+    = do Just zero <- shaped (hasArgs 0) cs
+              | Nothing => pure False
+         Just succ <- shaped (hasArgs 1) cs
+              | Nothing => pure False
+         let Just succCon = find (\con => name con == succ) cs
+              | Nothing => pure False
+         let Just (Evidence _ succArgTy) = firstArg (type succCon)
+              | Nothing => pure False
+         let Just succArgCon = typeCon succArgTy
+              | Nothing => pure False
+         if succArgCon == tyCon
+            then do setFlag fc zero (ConType ZERO)
+                    setFlag fc succ (ConType SUCC)
+                    pure True
+            else pure False
+calcNaty _ _ _ = pure False
+
+
 calcConInfo : {auto c : Ref Ctxt Defs} ->
-              FC -> List Constructor -> Core ()
-calcConInfo fc cons
-   = do False <- calcListy fc cons
+              FC -> Name -> List Constructor -> Core ()
+calcConInfo fc type cons
+   = do False <- calcNaty fc type cons
+           | True => pure ()
+        False <- calcListy fc cons
           | True => pure ()
        False <- calcMaybe fc cons
           | True => pure ()
@ -455,9 +492,9 @@ processData {vars} eopts nest env fc vis (MkImpData dfc n_in ty_raw opts cons_ra
         addToSave n
         log "declare.data" 10 $ "Saving from " ++ show n ++ ": " ++ show (keys (getMetas ty))

-         let connames = map conName cons
+         let connames = map name cons
         unless (NoHints `elem` opts) $
              traverse_ (\x => addHintFor fc (Resolved tidx) x True False) connames

-         calcConInfo fc cons
+         calcConInfo fc (Resolved tidx) cons
         traverse_ updateErasable (Resolved tidx :: connames)
--- a/tests/idris2/builtin006/Test.idr
+++ b/tests/idris2/builtin006/Test.idr
@ -1,9 +1,9 @@
 data MyNat
    = Z
-    | S MyNat
+    | S MyNat MyNat

 natToInt : MyNat -> Integer
 natToInt Z = 0
-natToInt (S k) = 1 + natToInt k
+natToInt (S k _) = 1 + natToInt k

 %builtin NaturalToInteger natToInt
--- a/tests/idris2/builtin006/expected
+++ b/tests/idris2/builtin006/expected
@ -4,7 +4,7 @@ Error: Non-erased argument is not a 'Nat'-like type.
 Test:9:1--9:35
 5 | natToInt : MyNat -> Integer
 6 | natToInt Z = 0
- 7 | natToInt (S k) = 1 + natToInt k
+ 7 | natToInt (S k _) = 1 + natToInt k
 8 | 
 9 | %builtin NaturalToInteger natToInt
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
--- a/tests/idris2/builtin008/expected
+++ b/tests/idris2/builtin008/expected
@ -1,5 +1,5 @@
 1/1: Building Test (Test.idr)
-Error: More than 1 non-erased arguments found for Main.finToInteger.
+Error: More than 1 non-erased arguments found for finToInteger.

 Test:11:1--11:39
 07 | finToInteger : {k : _} -> Fin k -> Integer
--- a/tests/idris2/builtin010/expected
+++ b/tests/idris2/builtin010/expected
@ -1,5 +1,5 @@
 1/1: Building Test (Test.idr)
-Error: No unrestricted arguments of type `Integer` found for Main.natToMyNat.
+Error: No unrestricted arguments of type `Integer` found for natToMyNat.

 Test:20:1--20:37
 16 | natToMyNat : Nat -> MyNat
--- a/tests/idris2/builtin012/Issue1799.idr
+++ b/tests/idris2/builtin012/Issue1799.idr
@ -0,0 +1,11 @@
+module Main
+
+import Data.Fin
+
+%builtin Natural Nat
+%builtin Natural Fin
+
+%builtin NaturalToInteger Data.Fin.finToInteger
+
+main : IO ()
+main = putStrLn $ show $ finToInteger $ FS $ FZ {k=34}
--- a/tests/idris2/builtin012/expected
+++ b/tests/idris2/builtin012/expected
@ -0,0 +1 @@
+1
--- a/tests/idris2/builtin012/run
+++ b/tests/idris2/builtin012/run
@ -0,0 +1,4 @@
+rm -rf build
+
+$1 --no-color --console-width 0 --no-banner --exec main Issue1799.idr
+