Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove nat magic #3184

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 33 additions & 17 deletions libs/prelude/Prelude/Types.idr
Original file line number Diff line number Diff line change
Expand Up @@ -30,17 +30,34 @@ prim__integerToNat i
then believe_me i
else Z

||| Convert a non-negative integer to a `Nat`
public export
nonNegativeIntegerToNat :
(x : Integer) ->
{auto 0 prf : (x >= 0) === True} ->
Nat
nonNegativeIntegerToNat 0 = Z
nonNegativeIntegerToNat x =
let -- x >= 0 and x != 0
-- so x > 0 so x - 1 >= 0
prf = believe_me Refl
in S $ nonNegativeIntegerToNat {prf}
$ assert_smaller x $ x - 1

public export
integerToNat : Integer -> Nat
integerToNat 0 = Z -- Force evaluation and hence caching of x at compile time
integerToNat x
= if intToBool (prim__lte_Integer x 0)
then Z
else S (assert_total (integerToNat (prim__sub_Integer x 1)))
integerToNat x with (x >= 0) proof prf
integerToNat x | True = nonNegativeIntegerToNat x {prf}
integerToNat x | False = Z

-- %builtin IntegerToNatural Prelude.Types.integerToNat
public export
natToInteger : Nat -> Integer
natToInteger Z = 0
natToInteger (S k) = 1 + natToInteger k
-- integer (+) may be non-linear in second
-- argument

-- Define separately so we can spot the name when optimising Nats
||| Add two natural numbers.
||| @ x the number to case-split on
||| @ y the other number
Expand All @@ -49,6 +66,8 @@ plus : (x : Nat) -> (y : Nat) -> Nat
plus Z y = y
plus (S k) y = S (plus k y)

%transform "Nat/plus" plus x y = integerToNat $ natToInteger x + natToInteger y

||| Subtract natural numbers.
||| If the second number is larger than the first, return 0.
public export
Expand All @@ -57,51 +76,48 @@ minus Z right = Z
minus left Z = left
minus (S left) (S right) = minus left right

%transform "Nat/minus" minus x y = integerToNat $ natToInteger x - natToInteger y

||| Multiply natural numbers.
public export
mult : (x : Nat) -> Nat -> Nat
mult Z y = Z
mult (S k) y = plus y (mult k y)

%transform "Nat/mult" mult x y = integerToNat $ natToInteger x * natToInteger y

public export
Num Nat where
(+) = plus
(*) = mult

fromInteger x = integerToNat x

-- used for nat hack
public export
equalNat : (m, n : Nat) -> Bool
equalNat Z Z = True
equalNat (S j) (S k) = equalNat j k
equalNat _ _ = False

%transform "Nat/equalNat" equalNat x y = natToInteger x == natToInteger y

public export
Eq Nat where
(==) = equalNat

-- used for nat hack
public export
compareNat : (m, n : Nat) -> Ordering
compareNat Z Z = EQ
compareNat Z (S k) = LT
compareNat (S k) Z = GT
compareNat (S j) (S k) = compareNat j k

%transform "Nat/compareNat" compareNat x y = compare (natToInteger x) (natToInteger y)

public export
Ord Nat where
compare = compareNat

public export
natToInteger : Nat -> Integer
natToInteger Z = 0
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
Expand Down
83 changes: 7 additions & 76 deletions src/Compiler/CompileExpr.idr
Original file line number Diff line number Diff line change
Expand Up @@ -142,77 +142,11 @@ mkDropSubst i es rest (x :: xs)
then (vs ** Drop sub)
else (x :: vs ** Keep sub)

-- Rewrite applications of Nat-like constructors and functions to more optimal
-- versions using Integer

-- None of these should be hard coded, but we'll do it this way until we
-- have a more general approach to optimising data types!
-- NOTE: Make sure that names mentioned here are listed in 'natHackNames' in
-- Common.idr, so that they get compiled, as they won't be spotted by the
-- usual calls to 'getRefs'.
data Magic : Type where
MagicCCon : Name -> (arity : Nat) -> -- checks
(FC -> forall vars. Vect arity (CExp vars) -> CExp vars) -> -- translation
Magic
MagicCRef : Name -> (arity : Nat) -> -- checks
(FC -> FC -> forall vars. Vect arity (CExp vars) -> CExp vars) -> --translation
Magic

magic : List Magic -> CExp vars -> CExp vars
magic ms (CLam fc x exp) = CLam fc x (magic ms exp)
magic ms e = go ms e where

fire : Magic -> CExp vars -> Maybe (CExp vars)
fire (MagicCCon n arity f) (CCon fc n' _ _ es)
= do guard (n == n')
map (f fc) (toVect arity es)
fire (MagicCRef n arity f) (CApp fc (CRef fc' n') es)
= do guard (n == n')
map (f fc fc') (toVect arity es)
fire _ _ = Nothing

go : List Magic -> CExp vars -> CExp vars
go [] e = e
go (m :: ms) e = case fire m e of
Nothing => go ms e
Just e' => e'

%inline
magic__integerToNat : FC -> FC -> forall vars. Vect 1 (CExp vars) -> CExp vars
magic__integerToNat fc fc' [k]
= CApp fc (CRef fc' (NS typesNS (UN $ Basic "prim__integerToNat"))) [k]

magic__natMinus : FC -> FC -> forall vars. Vect 2 (CExp vars) -> CExp vars
magic__natMinus fc fc' [m, n]
= magic__integerToNat fc fc'
[COp fc (Sub IntegerType) [m, n]]

-- We don't reuse natMinus here because we assume that unsuc will only be called
-- on S-headed numbers so we do not need the truncating integerToNat call!
magic__natUnsuc : FC -> FC -> forall vars. Vect 1 (CExp vars) -> CExp vars
magic__natUnsuc fc fc' [m]
||| Unsucc of some positive natural number
natUnsucc : FC -> FC -> forall vars. Vect 1 (CExp vars) -> CExp vars
natUnsucc fc fc' [m]
= COp fc (Sub IntegerType) [m, CPrimVal fc (BI 1)]

-- TODO: next release remove this and use %builtin pragma
natHack : List Magic
natHack =
[ MagicCRef (NS typesNS (UN $ Basic "natToInteger")) 1 (\ _, _, [k] => k)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This makes Nat -> integer a no-op right? Is that optimization captured anymore once this code is removed?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes that optimisation is captured already by a different, much more general optimisation I added.

, MagicCRef (NS typesNS (UN $ Basic "integerToNat")) 1 magic__integerToNat
, MagicCRef (NS typesNS (UN $ Basic "plus")) 2
(\ fc, fc', [m,n] => COp fc (Add IntegerType) [m, n])
, MagicCRef (NS typesNS (UN $ Basic "mult")) 2
(\ fc, fc', [m,n] => COp fc (Mul IntegerType) [m, n])
, MagicCRef (NS typesNS (UN $ Basic "minus")) 2 magic__natMinus
, MagicCRef (NS typesNS (UN $ Basic "equalNat")) 2
(\ fc, fc', [m,n] => COp fc (EQ IntegerType) [m, n])
, MagicCRef (NS typesNS (UN $ Basic "compareNat")) 2
(\ fc, fc', [m,n] => CApp fc (CRef fc' (NS eqOrdNS (UN $ Basic "compareInteger"))) [m, n])
]

-- get all builtin transformations
builtinMagic : forall vars. CExp vars -> CExp vars
builtinMagic = magic natHack

data NextMN : Type where
newMN : {auto s : Ref NextMN Int} -> String -> Core Name
newMN base = do
Expand All @@ -227,7 +161,7 @@ natBranch _ = False

trySBranch : CExp vars -> CConAlt vars -> Maybe (CExp vars)
trySBranch n (MkConAlt nm SUCC _ [arg] sc)
= Just (CLet (getFC n) arg YesInline (magic__natUnsuc (getFC n) (getFC n) [n]) sc)
= Just (CLet (getFC n) arg YesInline (natUnsucc (getFC n) (getFC n) [n]) sc)
trySBranch _ _ = Nothing

tryZBranch : CConAlt vars -> Maybe (CExp vars)
Expand Down Expand Up @@ -370,13 +304,10 @@ toCExp n tm
f' <- toCExpTm n f
Arity a <- numArgs defs f
| NewTypeBy arity pos =>
do let res = applyNewType arity pos f' args'
pure $ builtinMagic res
pure $ applyNewType arity pos f' args'
| EraseArgs arity epos =>
do let res = eraseConArgs arity epos f' args'
pure $ builtinMagic res
let res = expandToArity a f' args'
pure $ builtinMagic res
pure $ eraseConArgs arity epos f' args'
pure $ expandToArity a f' args'

mutual
conCases : {vars : _} ->
Expand Down
Loading