Skip to content

Commit ca69df6

Browse files
Qiu233Qiu233
committed
chore: optimize for performance
1 parent cef5eac commit ca69df6

2 files changed

Lines changed: 122 additions & 92 deletions

File tree

‎Binary/Get.lean‎

Lines changed: 118 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
module
22

33
public import Binary.Basic
4-
meta import Lean
54

65
namespace Binary
76

@@ -11,15 +10,14 @@ public section
1110
def fail (msg : String) : Get α :=
1211
throw (.userError msg)
1312

14-
@[specialize]
1513
def many (p : Get α) : Get (Array α) := do
1614
let mut data := #[]
1715
repeat
1816
let some x ← optional p | break
1917
data := data.push x
2018
return data
2119

22-
@[inline, specialize]
20+
@[inline]
2321
def many1 (p : Get α) : Get (Array α) := do
2422
let first ← p
2523
let rest ← many p
@@ -42,9 +40,6 @@ def notFollowedBy (p : Get α) : Get Unit := fun d =>
4240
| .error _ _ => DecodeResult.success () d
4341
| .pending _ => DecodeResult.error (.userError "unexpected pending lookahead") d
4442

45-
-- TODO: refactor following definitions for performance
46-
47-
@[inline, specialize]
4843
def takeAtLeast (n : Nat) (p : Get α) : Get (Array α) := do
4944
let mut r := Array.emptyWithCapacity n
5045
repeat
@@ -57,39 +52,39 @@ def takeAtLeast (n : Nat) (p : Get α) : Get (Array α) := do
5752
return r
5853

5954
/-- inclusive -/
60-
@[inline, specialize]
55+
@[inline]
6156
def takeUpTo (n : Nat) (p : Get α) : Get (Array α) := do
62-
let mut r := #[]
63-
repeat
57+
let mut r := Array.emptyWithCapacity n
58+
while true do
6459
if r.size == n then break
6560
let some x ← optional p | break
6661
r := r.push x
6762
return r
6863

6964
/-- inclusive -/
70-
@[inline, specialize]
65+
@[inline]
7166
def take1UpTo (n : Nat) (p : Get α) : Get (Array α) := do
7267
let x ← p
73-
let mut r := #[x]
68+
let mut r := (Array.emptyWithCapacity n).push x
7469
repeat
7570
if r.size == n then break
7671
let some x ← optional p | break
7772
r := r.push x
7873
return r
7974

80-
@[inline, specialize]
75+
@[inline]
8176
def takeN (n : Nat) (p : Get α) : Get (Array α) := do
82-
let mut r := Array.emptyWithCapacity 0
77+
let mut r := Array.emptyWithCapacity n
8378
repeat
8479
if r.size == n then break
8580
let x ← p
8681
r := r.push x
8782
return r
8883

8984
/--inclusive on both sides -/
90-
@[inline, specialize]
85+
@[inline]
9186
def takeRange (min max : Nat) (p : Get α) : Get (Array α) := do
92-
let mut r := Array.emptyWithCapacity min
87+
let mut r := Array.emptyWithCapacity max
9388
repeat
9489
if r.size == min then break
9590
let x ← p
@@ -100,7 +95,7 @@ def takeRange (min max : Nat) (p : Get α) : Get (Array α) := do
10095
r := r.push x
10196
return r
10297

103-
@[inline, specialize]
98+
@[inline]
10499
def sepBy (x : Get α) (sep : Get Unit) : Get (Array α) := do
105100
let some l ← optional x | return #[]
106101
let mut t := #[l]
@@ -109,7 +104,7 @@ def sepBy (x : Get α) (sep : Get Unit) : Get (Array α) := do
109104
t := t.push v
110105
return t
111106

112-
@[inline, specialize]
107+
@[inline]
113108
def sepBy1 (x : Get α) (s : Get Unit) : Get (Array α) := do
114109
let l ← x
115110
let mut t := #[l]
@@ -118,20 +113,20 @@ def sepBy1 (x : Get α) (s : Get Unit) : Get (Array α) := do
118113
t := t.push v
119114
return t
120115

121-
@[inline, specialize]
116+
@[inline]
122117
def sepByUpTo (n : Nat) (x : Get α) (s : Get Unit) : Get (Array α) := do
123118
let some l ← optional x | return #[]
124-
let mut t := #[l]
119+
let mut t := (Array.emptyWithCapacity n).push l
125120
repeat
126121
if t.size ≥ n then break
127122
let some v ← optional (s *> x) | break
128123
t := t.push v
129124
return t
130125

131-
@[inline, specialize]
126+
@[inline]
132127
def sepBy1UpTo (n : Nat) (x : Get α) (s : Get Unit) : Get (Array α) := do
133128
let l ← x
134-
let mut t := #[l]
129+
let mut t := (Array.emptyWithCapacity n).push l
135130
repeat
136131
if t.size ≥ n then break
137132
let some v ← optional (s *> x) | break
@@ -182,72 +177,62 @@ end
182177

183178
namespace Primitive
184179

185-
variable {ω m} [Monad m] [STWorld ω m] [MonadLiftT (ST ω) m]
186-
187-
private meta def generate_prim (le : Bool) (unsigned : Bool) (type : Lean.TSyntax `ident) (size : Lean.TSyntax `num) : Lean.MacroM Lean.Command := do
188-
let len := size.getNat
189-
if len = 0 then
190-
Lean.Macro.throwErrorAt size "size cannot be 0"
191-
let newSize := Lean.TSyntax.mk <| size.raw.setArg 0 (size.raw[0].setAtomVal s!"{len - 1}")
192-
let d ← Lean.mkIdent <$> Lean.Macro.addMacroScope `d
193-
let d_offset ← `($(Lean.mkIdent `Decoder.offset) $d:ident)
194-
let d_data ← `($(Lean.mkIdent `Decoder.data) $d:ident)
195-
let d_data_size ← `($(Lean.mkIdent `ByteArray.size) ($(Lean.mkIdent `Decoder.data) $d:ident))
196-
let ns := List.range len
197-
let ts ← ns.mapM fun x => do
198-
let y ←
199-
if unsigned then
200-
`($(Lean.mkIdent `ByteArray.get) $d_data ($d_offset + $(Lean.Syntax.mkNatLit x):num))
201-
else
202-
`($(Lean.mkIdent `ByteArray.get) $d_data ($d_offset + $(Lean.Syntax.mkNatLit x):num) |>.toInt8)
203-
let y ←
204-
if unsigned then
205-
`($(Lean.mkIdent (Lean.Name.mkStr2 "UInt8" s!"to{type.getId.getString!}")) $y)
206-
else
207-
`($(Lean.mkIdent (Lean.Name.mkStr2 "Int8" s!"to{type.getId.getString!}")) $y)
208-
let shift := if le then x * 8 else (len - 1 - x) * 8
209-
`($y <<< $(Lean.Syntax.mkNatLit shift):num)
210-
let combined ←
211-
match ts with
212-
| [] => unreachable!
213-
| [x] => pure x
214-
| head :: tail => do
215-
tail.foldlM (init := head) fun (x : Lean.Term) y => do
216-
`($x ||| $y)
217-
let code ← `(command|
218-
@[always_inline]
219-
scoped instance : Decode $type where
220-
get $d:ident :=
221-
if h : $d_offset + $newSize:num < $d_data_size then
222-
let val := $combined
223-
DecodeResult.success val {$d with offset := $d_offset + $(Lean.Syntax.mkNatLit len):num}
224-
else
225-
DecodeResult.mkEOI d
226-
)
227-
return code
228-
229-
local syntax "prim_unsigned_le " ident num : command
230-
local syntax "prim_unsigned_be " ident num : command
231-
local syntax "prim_signed_le " ident num : command
232-
local syntax "prim_signed_be " ident num : command
233-
234-
local macro_rules
235-
| `(command| prim_unsigned_le $type $size) => generate_prim true true type size
236-
| `(command| prim_unsigned_be $type $size) => generate_prim false true type size
237-
| `(command| prim_signed_le $type $size) => generate_prim true false type size
238-
| `(command| prim_signed_be $type $size) => generate_prim false false type size
239-
240180
public section
241181

242182
namespace LE
243183

244-
prim_unsigned_le UInt16 2
245-
prim_unsigned_le UInt32 4
246-
prim_unsigned_le UInt64 8
184+
@[always_inline]
185+
scoped instance : Decode UInt16 where
186+
get d :=
187+
if h : d.offset + 1 < d.data.size then
188+
let val :=
189+
(d.data[d.offset + 0]).toUInt16 |||
190+
(d.data[d.offset + 1]).toUInt16 <<< 8
191+
DecodeResult.success val {d with offset := d.offset + 2}
192+
else
193+
DecodeResult.mkEOI d
247194

248-
prim_signed_le Int16 2
249-
prim_signed_le Int32 4
250-
prim_signed_le Int64 8
195+
@[always_inline]
196+
scoped instance : Decode UInt32 where
197+
get d :=
198+
if h : d.offset + 3 < d.data.size then
199+
let val :=
200+
(d.data.get (d.offset + 0)).toUInt32 |||
201+
(d.data.get (d.offset + 1)).toUInt32 <<< 8 |||
202+
(d.data.get (d.offset + 2)).toUInt32 <<< 16 |||
203+
(d.data.get (d.offset + 3)).toUInt32 <<< 24
204+
DecodeResult.success val {d with offset := d.offset + 4}
205+
else
206+
DecodeResult.mkEOI d
207+
208+
@[always_inline]
209+
scoped instance : Decode UInt64 where
210+
get d :=
211+
if h : d.offset + 7 < d.data.size then
212+
let val :=
213+
(d.data.get (d.offset + 0)).toUInt64 |||
214+
(d.data.get (d.offset + 1)).toUInt64 <<< 8 |||
215+
(d.data.get (d.offset + 2)).toUInt64 <<< 16 |||
216+
(d.data.get (d.offset + 3)).toUInt64 <<< 24 |||
217+
(d.data.get (d.offset + 4)).toUInt64 <<< 32 |||
218+
(d.data.get (d.offset + 5)).toUInt64 <<< 40 |||
219+
(d.data.get (d.offset + 6)).toUInt64 <<< 48 |||
220+
(d.data.get (d.offset + 7)).toUInt64 <<< 56
221+
DecodeResult.success val {d with offset := d.offset + 8}
222+
else
223+
DecodeResult.mkEOI d
224+
225+
@[always_inline]
226+
scoped instance : Decode Int16 where
227+
get := Int16.ofUInt16 <$> Decode.get (α := UInt16)
228+
229+
@[always_inline]
230+
scoped instance : Decode Int32 where
231+
get := Int32.ofUInt32 <$> Decode.get (α := UInt32)
232+
233+
@[always_inline]
234+
scoped instance : Decode Int64 where
235+
get := Int64.ofUInt64 <$> Decode.get (α := UInt64)
251236

252237
@[always_inline]
253238
scoped instance : Decode Float32 where
@@ -261,13 +246,58 @@ end LE
261246

262247
namespace BE
263248

264-
prim_unsigned_be UInt16 2
265-
prim_unsigned_be UInt32 4
266-
prim_unsigned_be UInt64 8
249+
@[always_inline]
250+
scoped instance : Decode UInt16 where
251+
get d :=
252+
if h : d.offset + 1 < d.data.size then
253+
let val :=
254+
(d.data.get (d.offset + 0)).toUInt16 <<< 8 |||
255+
(d.data.get (d.offset + 1)).toUInt16
256+
DecodeResult.success val {d with offset := d.offset + 2}
257+
else
258+
DecodeResult.mkEOI d
259+
260+
@[always_inline]
261+
scoped instance : Decode UInt32 where
262+
get d :=
263+
if h : d.offset + 3 < d.data.size then
264+
let val :=
265+
(d.data.get (d.offset + 0)).toUInt32 <<< 24 |||
266+
(d.data.get (d.offset + 1)).toUInt32 <<< 16 |||
267+
(d.data.get (d.offset + 2)).toUInt32 <<< 8 |||
268+
(d.data.get (d.offset + 3)).toUInt32
269+
DecodeResult.success val {d with offset := d.offset + 4}
270+
else
271+
DecodeResult.mkEOI d
272+
273+
@[always_inline]
274+
scoped instance : Decode UInt64 where
275+
get d :=
276+
if h : d.offset + 7 < d.data.size then
277+
let val :=
278+
(d.data.get (d.offset + 0)).toUInt64 <<< 56 |||
279+
(d.data.get (d.offset + 1)).toUInt64 <<< 48 |||
280+
(d.data.get (d.offset + 2)).toUInt64 <<< 40 |||
281+
(d.data.get (d.offset + 3)).toUInt64 <<< 32 |||
282+
(d.data.get (d.offset + 4)).toUInt64 <<< 24 |||
283+
(d.data.get (d.offset + 5)).toUInt64 <<< 16 |||
284+
(d.data.get (d.offset + 6)).toUInt64 <<< 8 |||
285+
(d.data.get (d.offset + 7)).toUInt64
286+
DecodeResult.success val {d with offset := d.offset + 8}
287+
else
288+
DecodeResult.mkEOI d
267289

268-
prim_signed_be Int16 2
269-
prim_signed_be Int32 4
270-
prim_signed_be Int64 8
290+
@[always_inline]
291+
scoped instance : Decode Int16 where
292+
get := Int16.ofUInt16 <$> Decode.get (α := UInt16)
293+
294+
@[always_inline]
295+
scoped instance : Decode Int32 where
296+
get := Int32.ofUInt32 <$> Decode.get (α := UInt32)
297+
298+
@[always_inline]
299+
scoped instance : Decode Int64 where
300+
get := Int64.ofUInt64 <$> Decode.get (α := UInt64)
271301

272302
@[always_inline]
273303
scoped instance : Decode Float32 where

‎Binary/UTF8/Get.lean‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ private def byteToChar (b : UInt8) : Char :=
1515
private def chars_to_string (xs : Array Char) : String :=
1616
String.ofList xs.toList
1717

18-
@[always_inline, specialize]
18+
@[noinline]
1919
def satisfy (p : Char → Bool) : Get Char := do
2020
let b ← pending (getThe UInt8)
2121
let b1 := UInt8.toUInt32 b
@@ -75,7 +75,7 @@ def satisfy (p : Char → Bool) : Get Char := do
7575
@[always_inline]
7676
def pchar (c : Char) : Get Char := satisfy (· == c)
7777

78-
@[always_inline]
78+
@[inline]
7979
def pstring (s : String) : Get String := do
8080
for c in s.toList do
8181
_ ← inline pchar c
@@ -87,10 +87,10 @@ def skipChar (c : Char) : Get Unit := pchar c *> pure ()
8787
@[always_inline]
8888
def skipString (s : String) : Get Unit := pstring s *> pure ()
8989

90-
@[always_inline, specialize]
90+
@[always_inline]
9191
def manyChars (p : Get Char) : Get String :=
9292
chars_to_string <$> many p
9393

94-
@[always_inline, specialize]
94+
@[always_inline]
9595
def many1Chars (p : Get Char) : Get String :=
9696
chars_to_string <$> many1 p

0 commit comments

Comments
 (0)