82 lines
2.6 KiB
Text
82 lines
2.6 KiB
Text
/-! Mixing basic IO actions and heavy unboxed compute -/
|
||
|
||
def N : Nat := 12500000 / 2 -- ~50 MB
|
||
|
||
-- xorshift64* PRNG (same algorithm as Rust version)
|
||
@[inline] def xorshift64 (state : UInt64) : UInt64 × UInt64 :=
|
||
let x := state ^^^ (state >>> 12)
|
||
let x := x ^^^ (x <<< 25)
|
||
let x := x ^^^ (x >>> 27)
|
||
(x, x * 0x2545F4914F6CDD1D)
|
||
|
||
-- Push UInt64 as 8 little-endian bytes
|
||
@[inline] def pushLE (ba : ByteArray) (v : UInt64) : ByteArray :=
|
||
let ba := ba.push (v &&& 0xFF).toUInt8
|
||
let ba := ba.push ((v >>> 8) &&& 0xFF).toUInt8
|
||
let ba := ba.push ((v >>> 16) &&& 0xFF).toUInt8
|
||
let ba := ba.push ((v >>> 24) &&& 0xFF).toUInt8
|
||
let ba := ba.push ((v >>> 32) &&& 0xFF).toUInt8
|
||
let ba := ba.push ((v >>> 40) &&& 0xFF).toUInt8
|
||
let ba := ba.push ((v >>> 48) &&& 0xFF).toUInt8
|
||
ba.push ((v >>> 56) &&& 0xFF).toUInt8
|
||
|
||
-- Read UInt64 from 8 little-endian bytes at offset
|
||
@[inline] def readLE (ba : ByteArray) (off : Nat) : UInt64 :=
|
||
(ba.get! off).toUInt64 +
|
||
((ba.get! (off + 1)).toUInt64 <<< 8) +
|
||
((ba.get! (off + 2)).toUInt64 <<< 16) +
|
||
((ba.get! (off + 3)).toUInt64 <<< 24) +
|
||
((ba.get! (off + 4)).toUInt64 <<< 32) +
|
||
((ba.get! (off + 5)).toUInt64 <<< 40) +
|
||
((ba.get! (off + 6)).toUInt64 <<< 48) +
|
||
((ba.get! (off + 7)).toUInt64 <<< 56)
|
||
|
||
def timeS {α : Type} (act : IO α) : IO (α × Float) := do
|
||
let t0 ← IO.monoMsNow
|
||
let r ← act
|
||
let t1 ← IO.monoMsNow
|
||
return (r, (t1 - t0).toFloat / 1000.0)
|
||
|
||
def main : IO Unit := do
|
||
-- Phase 1: Generate & Sort
|
||
let (data, elapsed) ← timeS do
|
||
let mut state : UInt64 := 42
|
||
let mut arr : Array UInt64 := Array.mkEmpty N
|
||
for _ in [:N] do
|
||
let (s, v) := xorshift64 state
|
||
state := s
|
||
arr := arr.push v
|
||
return arr.qsortOrd
|
||
IO.println s!"measurement: generate_sort {elapsed} s"
|
||
|
||
let (_, file) ← IO.FS.createTempFile
|
||
-- Phase 2: Write
|
||
let (_, elapsed) ← timeS do
|
||
let mut ba := ByteArray.emptyWithCapacity (N * 8)
|
||
for i in [:data.size] do
|
||
ba := pushLE ba data[i]!
|
||
IO.FS.writeBinFile file ba
|
||
IO.println s!"measurement: write {elapsed} s"
|
||
|
||
-- Phase 3: Read
|
||
let (data, elapsed) ← timeS do
|
||
let ba ← IO.FS.readBinFile file
|
||
let mut arr : Array UInt64 := Array.mkEmpty N
|
||
for i in [:N] do
|
||
arr := arr.push (readLE ba (i * 8))
|
||
return arr
|
||
IO.println s!"measurement: read {elapsed} s"
|
||
|
||
-- Phase 4: Fisher-Yates shuffle
|
||
let (_, elapsed) ← timeS do
|
||
let mut arr := data
|
||
let mut state : UInt64 := 42
|
||
for i' in [:N - 1] do
|
||
let i := N - 1 - i'
|
||
let (s, v) := xorshift64 state
|
||
state := s
|
||
let j := (v % (i + 1).toUInt64).toNat
|
||
arr := arr.swapIfInBounds i j
|
||
return arr
|
||
IO.println s!"measurement: shuffle {elapsed} s"
|
||
|