leanexe/LeanExe/Examples/IntMap.lean at main · jsmorph/leanexe · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
namespace LeanExe.Examples.IntMap

def capacityNat : Nat :=
  256

def capacity : UInt64 :=
  256

structure Slot where
  key : UInt64
  value : UInt64
  deriving Inhabited

structure Table where
  slots : Array Slot

def emptySlot : Slot :=
  { key := 0, value := 0 }

def empty : Table :=
  { slots := Array.replicate capacityNat emptySlot }

def hash (key : UInt64) : UInt64 :=
  key % capacity

def probe (key attempt : UInt64) : UInt64 :=
  (hash key + attempt) % capacity

def getSlot (table : Table) (slot : UInt64) : Slot :=
  table.slots[slot.toNat]!

def getKey (table : Table) (slot : UInt64) : UInt64 :=
  (getSlot table slot).key

def getValue (table : Table) (slot : UInt64) : UInt64 :=
  (getSlot table slot).value

def setSlot (table : Table) (slot key value : UInt64) : Table :=
  { slots := table.slots.set! slot.toNat { key := key, value := value } }

def slotAvailable (table : Table) (key attempt : UInt64) : Bool :=
  getKey table (probe key attempt) == 0 || getKey table (probe key attempt) == key

def shouldInsert (table : Table) (key attempt done : UInt64) : Bool :=
  done == 0 && slotAvailable table key attempt

def insertFuel : Nat → UInt64 → UInt64 → UInt64 → UInt64 → Table → Table
  | 0, _, _, _, _, table => table
  | fuel + 1, key, value, attempt, done, table =>
      insertFuel fuel key value (attempt + 1)
        (if shouldInsert table key attempt done then 1 else done)
        (if shouldInsert table key attempt done then
          setSlot table (probe key attempt) key value
        else
          table)

def insert (key value : UInt64) (table : Table) : Table :=
  insertFuel capacityNat key value 0 0 table

def slotMatches (table : Table) (key attempt found : UInt64) : Bool :=
  found == 0 && getKey table (probe key attempt) == key

def lookupFuel : Nat → UInt64 → UInt64 → UInt64 → Table → UInt64 → UInt64
  | 0, _, _, _, _, result => result
  | fuel + 1, key, attempt, found, table, result =>
      lookupFuel fuel key (attempt + 1)
        (if slotMatches table key attempt found then 1 else found)
        table
        (if slotMatches table key attempt found then
          getValue table (probe key attempt)
        else
          result)

def lookup (key : UInt64) (table : Table) : UInt64 :=
  lookupFuel capacityNat key 0 0 table 0

def buildFuel : Nat → UInt64 → Table → Table
  | 0, _, table => table
  | fuel + 1, key, table =>
      buildFuel fuel (key + 1) (insert key (key * 10 + 7) table)

def build : Table :=
  buildFuel 100 1 empty

def checksumFuel : Nat → UInt64 → Table → UInt64 → UInt64
  | 0, _, _, acc => acc
  | fuel + 1, key, table, acc =>
      checksumFuel fuel (key + 1) table (acc + lookup key table)

def checksum : UInt64 :=
  checksumFuel 100 1 build 0

def query (key : UInt64) : UInt64 :=
  lookup key build

end LeanExe.Examples.IntMap