1. 1
2. 2
3. 3
4. 4
5. 5
6. 6
7. 7
8. 8
9. 9
10. 10
11. 11
12. 12
13. 13
14. 14
15. 15
16. 16
17. 17
18. 18
19. 19
20. 20
21. 21
22. 22
23. 23
24. 24
25. 25
26. 26
27. 27
28. 28
29. 29
30. 30
31. 31
32. 32
33. 33
34. 34
35. 35
36. 36
37. 37
38. 38
39. 39
40. 40
41. 41
42. 42
43. 43
44. 44
45. 45
46. 46
47. 47
48. 48
49. 49
50. 50
51. 51
52. 52
53. 53
54. 54
55. 55
56. 56
57. 57
58. 58
59. 59
60. 60
61. 61
62. 62
63. 63
64. 64
65. 65
66. 66
67. 67
68. 68
69. 69
70. 70
71. 71
72. 72
73. 73
74. 74
75. 75
76. 76
77. 77
78. 78
79. 79
80. 80
81. 81
82. 82
83. 83
84. 84
85. 85
86. 86
87. 87
88. 88
89. 89
90. 90
91. 91
92. 92
93. 93
94. 94
95. 95
96. 96
97. 97
98. 98
99. 99
100. 100
101. 101
102. 102
103. 103
104. 104
105. 105
106. 106
107. 107
108. 108
109. 109
110. 110
111. 111
112. 112
113. 113
114. 114
115. 115
116. 116
117. 117
118. 118
119. 119
120. 120
121. 121
122. 122
123. 123
124. 124
125. 125
126. 126
127. 127
128. 128
129. 129
130. 130
131. 131
132. 132
133. 133
134. 134
135. 135
136. 136
137. 137
138. 138
139. 139
140. 140
141. 141
142. 142
143. 143
144. 144
145. 145
146. 146
147. 147
148. 148
149. 149
150. 150
151. 151
152. 152
153. 153
154. 154
155. 155
156. 156
157. 157
158. 158
159. 159
160. 160
161. 161
162. 162
163. 163
164. 164
165. 165
166. 166
167. 167
168. 168
169. 169
170. 170
171. 171
172. 172
173. 173
174. 174
175. 175
176. 176
177. 177
178. 178
179. 179
180. 180
181. 181

import Std.Data

def TypeListToTuple : List Type → Type
  | [] => Unit
  | [α] => α
  | α :: αs => α × TypeListToTuple αs

inductive DTerm (vars : String → Type) (α : Type)
  | const (x : α)
  | var (s : String) (h : vars s = α := by rfl)

def TypeListToDTerms (vars : String → Type) : List Type → Type
  | [] => Unit
  | [α] => DTerm vars α
  | α :: αs => DTerm vars α × TypeListToDTerms vars αs

def Atom (relations : String → List Type) (vars : String → Type) :=
  (rel : String) × TypeListToDTerms vars (relations rel)

def Rule relations vars :=
  Atom relations vars × List (Atom relations vars)

def unify
  (αs : List Type)
  (val : TypeListToTuple αs)
  (ds : TypeListToDTerms vars αs)
  (env : Std.DHashMap String (vars ·))
  (h : ∀ α ∈ αs, BEq α)
  : Option (Std.DHashMap String (vars ·)) :=
  match αs with
  | [] => some env
  | [α] => by
    have := h α (by grind)
    simp [TypeListToTuple] at val
    simp [TypeListToDTerms] at ds
    match ds with
    | .const x =>
      exact if val == x then
        some env
      else
        none
    | .var s h =>
      exact if he : s ∈ env then
        if val == h ▸ env.get s he then
          some env
        else
          none
      else
        some <| env.insert s (h ▸ val)
  | α :: α' :: αs => by
    have := h α (by grind)
    simp [TypeListToTuple] at val
    simp [TypeListToDTerms] at ds
    match ds.1 with
    | .const x =>
      exact if val.1 == x then
        unify (α' :: αs) val.2 ds.2 env (by solve_by_elim)
      else
        none
    | .var s h =>
      exact if he : s ∈ env then
        if val.1 == h ▸ env.get s he then
          unify (α' :: αs) val.2 ds.2 env (by solve_by_elim)
        else
          none
      else
        unify (α' :: αs) val.2 ds.2 (env.insert s (h ▸ val.1)) (by solve_by_elim)

def headToTuple
  (rel : String)
  (αs : List Type)
  (ds : TypeListToDTerms vars αs)
  (env : Std.DHashMap String (vars ·))
  (hi : ∀ s, Inhabited (vars s))
  : TypeListToTuple αs :=
  match αs with
  | [] => ()
  | [α] => by
    simp [TypeListToDTerms] at ds
    simp [TypeListToTuple]
    match ds with
    | .const x => exact x
    | .var s h => exact h ▸ env.get! s
  | α :: α' :: αs => by
    simp [TypeListToDTerms] at ds
    simp [TypeListToTuple]
    let ret := headToTuple rel (α' :: αs) ds.2 env hi
    match ds.1 with
    | .const x => exact (x, ret)
    | .var s h => exact (h ▸ env.get! s, ret)

def naiveEval
  (rules : List (Rule relations vars))
  (hb : ∀ rel, BEq (TypeListToTuple <| relations rel))
  (hh : ∀ rel, Hashable (TypeListToTuple <| relations rel))
  (hi : ∀ s, Inhabited (vars s))
  (hb' : ∀ rel, ∀ α ∈ relations rel, BEq α)
  : Std.DHashMap String (Std.HashSet <| TypeListToTuple <| relations ·) := Id.run do
  let RetType := Std.DHashMap String (Std.HashSet <| TypeListToTuple <| relations ·)
  let mut cur := .ofList []
  repeat
    let mut changed := false
    let mut nxt := cur
    for rule in rules do
      let rec unifyRec (env : Std.DHashMap String (vars ·)) (changed : Bool) (cur nxt : RetType) : List (Atom relations vars) → Bool × RetType
        | [] =>
          let tup := headToTuple rule.1.1 (relations rule.1.1) rule.1.2 env hi
          if rule.1.1 ∈ nxt then
            if (if h : rule.1.1 ∈ cur then tup ∈ cur.get rule.1.1 h else false) then
              (changed, nxt)
            else
              (true, nxt.modify rule.1.1 (·.insert tup))
          else
            (true, nxt.insert rule.1.1 (.ofList [tup]))
        | atom :: atoms => Id.run do
          let mut changed := changed
          let mut nxt := nxt
          if h : atom.1 ∈ cur then
            for val in cur.get atom.1 h do
              match unify (relations atom.1) val atom.2 env (hb' atom.1) with
              | some env =>
                (changed, nxt) := unifyRec env changed cur nxt atoms
              | none =>
                ()
          return (changed, nxt)
      (changed, nxt) := unifyRec (.ofList []) changed cur nxt rule.2
    if !changed then
      break
    cur := nxt
  return cur

def relations : String → List Type
  | "edge" | "path" => [Nat, Nat]
  | _ => []

def vars : String → Type := fun _ ↦ Nat

def pathRule : Rule relations vars :=
  (⟨"path", .var "x", .var "y"⟩, [⟨"edge", .var "x", .var "y"⟩])

def extendPathRule : Rule relations vars :=
  (⟨"path", .var "x", .var "z"⟩, [
    ⟨"path", .var "x", .var "y"⟩,
    ⟨"edge", .var "y", .var "z"⟩
  ])

def main := do
  let rules := [
    pathRule,
    extendPathRule,
    (⟨"edge", .const 1, .const 2⟩, []),
    (⟨"edge", .const 2, .const 3⟩, []),
    (⟨"edge", .const 2, .const 4⟩, []),
    (⟨"edge", .const 4, .const 1⟩, []),
  ]
  -- The default BEq instance produced by solve_by_elim is bad
  have relsBEq (rel : String) : BEq (TypeListToTuple (relations rel)) := by
    by_cases h : rel = "edge"
    · simpa [h, relations, TypeListToTuple] using instBEqProd
    · by_cases h : rel = "path"
      · simpa [h, relations, TypeListToTuple] using instBEqProd
      · simp [relations, TypeListToTuple]
        exact instBEqOfDecidableEq
  have relsHashable (rel : String) : Hashable (TypeListToTuple (relations rel)) := by
    by_cases h : rel = "edge"
    · simpa [h, relations, TypeListToTuple] using instHashableProd
    · by_cases h : rel = "path"
      · simpa [h, relations, TypeListToTuple] using instHashableProd
      · simpa [relations, TypeListToTuple] using instHashablePUnit
  have relsBEq' (rel : String) α (h : α ∈ relations rel) : BEq α := by
    by_cases hr : rel = "edge"
    · simp [hr, relations] at h
      simpa [h] using instBEqOfDecidableEq
    · by_cases hr : rel = "path"
      · simp [hr, relations] at h
        simpa [h] using instBEqOfDecidableEq
      · simp [relations] at h
  let ans := naiveEval rules relsBEq relsHashable (by solve_by_elim) relsBEq'
  have : TypeListToTuple [Nat × Nat] = (Nat × Nat) := by rfl
  IO.println <| this ▸ (ans.get! "edge" |>.toList)
  IO.println <| this ▸ (ans.get! "path" |>.toList)