Move AStar into a separte library and use it for Day17-1

1 files changed, 57 insertions, 128 deletions

diff --git a/Day17.lean b/Day17.lean
index 11a92f2..769474b 100644
--- a/Day17.lean
+++ b/Day17.lean
@@ -1,6 +1,7 @@
 import Common
 import Std.Data.HashSet
 import BinaryHeap
+import LeanAStar
 
 namespace Day17
 
@@ -91,9 +92,9 @@ private def StepsInDirection.next (s : StepsInDirection) (h₁ : s ≠ .Three) :
 
 private structure PathNode (heatLossMap : HeatLossMap) where
   coordinate : heatLossMap.Coordinate
-  accumulatedCosts : Nat
   currentDirection : Direction
   takenSteps : StepsInDirection
+deriving BEq, Hashable
 
 private def PathNode.goUp? (node : PathNode heatLossMap) : Option (PathNode heatLossMap) :=
   match node.coordinate.y, node.currentDirection, node.takenSteps with
@@ -106,7 +107,6 @@ private def PathNode.goUp? (node : PathNode heatLossMap) : Option (PathNode heat
     let coordinate := {x := node.coordinate.x, y := ⟨y, Nat.lt_of_succ_lt h₁⟩}
     some {
       coordinate,
-      accumulatedCosts := node.accumulatedCosts + heatLossMap[coordinate],
       currentDirection := .Up,
       takenSteps,
       }
@@ -114,7 +114,6 @@ private def PathNode.goUp? (node : PathNode heatLossMap) : Option (PathNode heat
     let coordinate := { x := node.coordinate.x, y := ⟨y, Nat.lt_of_succ_lt h₁⟩}
     some {
       coordinate,
-      accumulatedCosts := node.accumulatedCosts + heatLossMap[coordinate],
       currentDirection := .Up,
       takenSteps := .One,
     }
@@ -138,7 +137,6 @@ private def PathNode.goLeft? (node : PathNode heatLossMap) : Option (PathNode he
     let coordinate := { x := ⟨x, Nat.lt_of_succ_lt h₁⟩, y := node.coordinate.y }
     some {
       coordinate,
-      accumulatedCosts := node.accumulatedCosts + heatLossMap[coordinate],
       currentDirection := .Left
       takenSteps
     }
@@ -146,7 +144,6 @@ private def PathNode.goLeft? (node : PathNode heatLossMap) : Option (PathNode he
     let coordinate := { x := ⟨x, Nat.lt_of_succ_lt h₁⟩, y := node.coordinate.y }
     some {
       coordinate,
-      accumulatedCosts := node.accumulatedCosts + heatLossMap[coordinate],
       currentDirection := .Left
       takenSteps := .One
     }
@@ -170,7 +167,6 @@ private def PathNode.goDown? (node : PathNode heatLossMap) : Option (PathNode he
     let coordinate := {x := node.coordinate.x, y := Fin.rev ⟨y, Nat.lt_of_succ_lt h₁⟩}
     some {
       coordinate,
-      accumulatedCosts := node.accumulatedCosts + heatLossMap[coordinate],
       currentDirection := .Down,
       takenSteps,
       }
@@ -178,7 +174,6 @@ private def PathNode.goDown? (node : PathNode heatLossMap) : Option (PathNode he
     let coordinate := { x := node.coordinate.x, y := Fin.rev ⟨y, Nat.lt_of_succ_lt h₁⟩}
     some {
       coordinate,
-      accumulatedCosts := node.accumulatedCosts + heatLossMap[coordinate],
       currentDirection := .Down,
       takenSteps := .One,
     }
@@ -205,7 +200,6 @@ private def PathNode.goRight? (node : PathNode heatLossMap) : Option (PathNode h
     let coordinate := {x := Fin.rev ⟨x, Nat.lt_of_succ_lt h₁⟩, y := node.coordinate.y}
     some {
       coordinate,
-      accumulatedCosts := node.accumulatedCosts + heatLossMap[coordinate],
       currentDirection := .Right,
       takenSteps,
       }
@@ -213,7 +207,6 @@ private def PathNode.goRight? (node : PathNode heatLossMap) : Option (PathNode h
     let coordinate := {x := Fin.rev ⟨x, Nat.lt_of_succ_lt h₁⟩, y := node.coordinate.y}
     some {
       coordinate,
-      accumulatedCosts := node.accumulatedCosts + heatLossMap[coordinate],
       currentDirection := .Right,
       takenSteps := .One,
     }
@@ -242,20 +235,6 @@ private def PathNode.estimateMinimumCostToGoal (node : PathNode heatLossMap) : N
   }
   (goal.x - node.coordinate.x : Fin _).val + (goal.y - node.coordinate.y : Fin _).val
 
-private def PathNode.heuristics (node : PathNode heatLossMap) : Nat :=
-  node.estimateMinimumCostToGoal + node.accumulatedCosts
-
-private def PathNode.heuristicsLe (a b : PathNode heatLossMap) : Bool :=
-  Nat.ble a.heuristics b.heuristics
-
-theorem PathNode.heuristicsLe_transitive : BinaryHeap.transitive_le (PathNode.heuristicsLe (heatLossMap := heatLossMap)) :=
-  λa b c ↦ BinaryHeap.nat_ble_to_heap_transitive_le a.heuristics b.heuristics c.heuristics
-
-theorem PathNode.heuristicsLe_total : BinaryHeap.total_le (PathNode.heuristicsLe (heatLossMap := heatLossMap)) :=
-  λa b ↦ BinaryHeap.nat_ble_to_heap_le_total a.heuristics b.heuristics
-
-private theorem PathNode.heuristicsLe_total_and_transitive : BinaryHeap.TotalAndTransitiveLe (PathNode.heuristicsLe (heatLossMap := heatLossMap)) := ⟨PathNode.heuristicsLe_transitive, PathNode.heuristicsLe_total⟩
-
 private def PathNode.isGoal (node : PathNode heatLossMap) : Bool :=
   let goal : heatLossMap.Coordinate := {
     x := ⟨heatLossMap.width - 1, Nat.pred_lt_self heatLossMap.not_empty.left⟩,
@@ -263,134 +242,81 @@ private def PathNode.isGoal (node : PathNode heatLossMap) : Bool :=
   }
   node.coordinate == goal
 
-end PathNode
-
-abbrev OpenSet (heatLossMap : HeatLossMap) := BinaryHeap (PathNode heatLossMap) PathNode.heuristicsLe
-
-private def HeatLossMap.start (heatLossMap : HeatLossMap) : List (PathNode heatLossMap) :=
-  let a : List (PathNode heatLossMap) := if h : heatLossMap.width > 1 then
-    let coordinate := {x := ⟨1,h⟩, y := ⟨0, heatLossMap.not_empty.right⟩ }
-    [{
-      coordinate,
-      accumulatedCosts := heatLossMap[coordinate],
-      currentDirection := .Right,
-      takenSteps := .One
-    }]
-  else
-    []
-  if h : heatLossMap.height > 1 then
-    let coordinate := {x := ⟨0, heatLossMap.not_empty.left⟩, y := ⟨1,h⟩}
-    {
-      coordinate,
-      accumulatedCosts := heatLossMap[coordinate]
-      currentDirection := .Down,
-      takenSteps := .One
-    } :: a
-  else
-    a
 
-private def OpenSet.start (heatLossMap : HeatLossMap) : OpenSet heatLossMap (heatLossMap.start.length) :=
-  --we cannot add the start tile directly - it's invalid state, as it doesn't have a direction
-  BinaryHeap.ofList PathNode.heuristicsLe_total_and_transitive heatLossMap.start
-
-private structure ClosedSetEntry (heatLossMap : HeatLossMap) where
-  coordinate : heatLossMap.Coordinate
-  direction : Direction
-  steps : StepsInDirection
-deriving BEq, Hashable
-
-instance {heatLossMap : HeatLossMap} : LawfulBEq (ClosedSetEntry heatLossMap) where
+instance : LawfulBEq (PathNode heatLossMap) where
   rfl := λ {a} ↦
     match a with
-    | {coordinate, direction, steps} => by
-      unfold BEq.beq instBEqClosedSetEntry
+    | {coordinate, currentDirection, takenSteps} => by
+      unfold BEq.beq instBEqPathNode
       simp! --no clue how to rename an unnamed function in the goal
   eq_of_beq := λ{a b} h₁ ↦ by
-    unfold BEq.beq instBEqClosedSetEntry at h₁
+    unfold BEq.beq instBEqPathNode at h₁
     cases a
     cases b
     simp! at h₁
     simp[h₁]
 
-private def ClosedSetEntry.toTuple {heatLossMap : HeatLossMap} : (ClosedSetEntry heatLossMap) → (heatLossMap.Coordinate × Direction × StepsInDirection)
-| {coordinate, direction, steps} => (coordinate, direction, steps)
+private def PathNode.toTuple : (PathNode heatLossMap) → (heatLossMap.Coordinate × Direction × StepsInDirection)
+| {coordinate, currentDirection, takenSteps} => (coordinate, currentDirection, takenSteps)
 
-private def ClosedSetEntry.ofTuple {heatLossMap : HeatLossMap} : (heatLossMap.Coordinate × Direction × StepsInDirection) → (ClosedSetEntry heatLossMap)
-| (coordinate, direction, steps) => {coordinate, direction, steps}
+private def PathNode.ofTuple  : (heatLossMap.Coordinate × Direction × StepsInDirection) → (PathNode heatLossMap)
+| (coordinate, currentDirection, takenSteps) => {coordinate, currentDirection, takenSteps}
 
-private theorem ClosedSetEntry.toTuple_inv_ofTuple {heatLossMap : HeatLossMap} : (ClosedSetEntry.toTuple (heatLossMap := heatLossMap)) ∘ ClosedSetEntry.ofTuple = id := rfl
-private theorem ClosedSetEntry.ofTuple_inv_toTuple {heatLossMap : HeatLossMap} : (ClosedSetEntry.ofTuple (heatLossMap := heatLossMap)) ∘ ClosedSetEntry.toTuple = id := rfl
+private theorem PathNode.toTuple_inv_ofTuple : (PathNode.toTuple (heatLossMap := heatLossMap)) ∘ PathNode.ofTuple = id := rfl
+private theorem PathNode.ofTuple_inv_toTuple : (PathNode.ofTuple (heatLossMap := heatLossMap)) ∘ PathNode.toTuple = id := rfl
 
-instance {heatLossMap : HeatLossMap} : Finite (ClosedSetEntry heatLossMap) where
+instance : Finite (PathNode heatLossMap) where
   cardinality := Finite.cardinality (heatLossMap.Coordinate × Direction × StepsInDirection)
-  enumerate := Finite.enumerate ∘ ClosedSetEntry.toTuple
-  nth := ClosedSetEntry.ofTuple ∘ Finite.nth
+  enumerate := Finite.enumerate ∘ PathNode.toTuple
+  nth := PathNode.ofTuple ∘ Finite.nth
   enumerate_inverse_nth := by
     funext x
     rewrite[Function.comp_assoc]
     rewrite (occs := .pos [2]) [←Function.comp_assoc]
-    simp only[ClosedSetEntry.toTuple_inv_ofTuple, Function.id_comp, Finite.enumerate_inverse_nth]
+    simp only[PathNode.toTuple_inv_ofTuple, Function.id_comp, Finite.enumerate_inverse_nth]
   nth_inverse_enumerate := by
     funext x
     rewrite[Function.comp_assoc]
     rewrite (occs := .pos [2]) [←Function.comp_assoc]
-    simp only[Finite.nth_inverse_enumerate, Function.id_comp, ClosedSetEntry.ofTuple_inv_toTuple]
-
-instance {heatLossMap : HeatLossMap} : Coe (PathNode heatLossMap) (ClosedSetEntry heatLossMap) where
-  coe := λ{coordinate, currentDirection, takenSteps, ..} ↦ {coordinate, direction := currentDirection, steps := takenSteps}
-
-abbrev ClosedSet (heatLossMap : HeatLossMap) := Std.HashSet (ClosedSetEntry heatLossMap)
-
-private def OpenSet.findFirstNotInClosedSet {heatLossMap : HeatLossMap} {n : Nat} (openSet : OpenSet heatLossMap n) (closedSet : ClosedSet heatLossMap) : Option ((r : Nat) × PathNode heatLossMap × OpenSet heatLossMap r) :=
-  match n, openSet with
-  | 0, _ => none
-  | m+1, openSet =>
-    let (node, openSet) := openSet.pop
-    if closedSet.contains node then
-      findFirstNotInClosedSet openSet closedSet
-    else
-      some ⟨m, node, openSet⟩
-
-private theorem OpenSet.findFirstNotInClosedSet_not_in_closed_set {heatLossMap : HeatLossMap} {n : Nat} (openSet : OpenSet heatLossMap n) (closedSet : ClosedSet heatLossMap) {result : (r : Nat) × PathNode heatLossMap × OpenSet heatLossMap r} (h₁ : openSet.findFirstNotInClosedSet closedSet = some result) : ¬closedSet.contains result.snd.fst := by
-  simp
-  unfold findFirstNotInClosedSet at h₁
-  split at h₁; contradiction
-  simp at h₁
-  split at h₁
-  case h_2.isTrue =>
-    have h₃ := findFirstNotInClosedSet_not_in_closed_set _ closedSet h₁
-    simp at h₃
-    assumption
-  case h_2.isFalse h₂ =>
-    simp at h₂ h₁
-    subst result
-    assumption
-
-private def HeatLossMap.findPath {heatLossMap : HeatLossMap} {n : Nat} (openSet : OpenSet heatLossMap n) (closedSet : ClosedSet heatLossMap) : Option Nat :=
-  match h₁ : openSet.findFirstNotInClosedSet closedSet with
-  | none => none
-  | some ⟨_,(node, openSet)⟩ =>
-    if node.isGoal then
-      some node.accumulatedCosts
-    else
-      let neighbours := node.getNeighbours.filter (not ∘ closedSet.contains ∘ Coe.coe)
-      let newClosedSet := closedSet.insert node
-      let openSet := openSet.pushList neighbours
-      findPath openSet newClosedSet
-termination_by (Finite.cardinality (ClosedSetEntry heatLossMap)) - closedSet.size
-decreasing_by
-  have h₃ := Std.HashSet.size_insert (m := closedSet) (k := { coordinate := node.coordinate, direction := node.currentDirection, steps := node.takenSteps })
-  split at h₃
-  case isTrue =>
-    have := OpenSet.findFirstNotInClosedSet_not_in_closed_set _ _ h₁
-    contradiction
-  case isFalse h₂ =>
-    rw[h₃]
-    have : closedSet.size < (Finite.cardinality (ClosedSetEntry heatLossMap)) := Std.HashSet.size_lt_finite_cardinality_of_not_mem closedSet ⟨_,h₂⟩
-    omega
+    simp only[Finite.nth_inverse_enumerate, Function.id_comp, PathNode.ofTuple_inv_toTuple]
+
+instance : LeanAStar.AStarNode (PathNode heatLossMap) where
+  Costs := Nat
+  costsLe := Nat.ble
+  costs_order := ⟨BinaryHeap.nat_ble_to_heap_transitive_le, BinaryHeap.nat_ble_to_heap_le_total⟩
+  remaining_costs_heuristic := PathNode.estimateMinimumCostToGoal
+  isGoal := PathNode.isGoal
+  getNeighbours := λn↦
+    let ns := n.getNeighbours
+    ns.map λn ↦ (n, heatLossMap[n.coordinate])
+end PathNode
+
+private def HeatLossMap.start (heatLossMap : HeatLossMap) : List (PathNode heatLossMap) :=
+  let a : List (PathNode heatLossMap) := if h : heatLossMap.width > 1 then
+    let coordinate := {x := ⟨1,h⟩, y := ⟨0, heatLossMap.not_empty.right⟩ }
+    [{
+      coordinate,
+      currentDirection := .Right,
+      takenSteps := .One
+    }]
+  else
+    []
+  if h : heatLossMap.height > 1 then
+    let coordinate := {x := ⟨0, heatLossMap.not_empty.left⟩, y := ⟨1,h⟩}
+    {
+      coordinate,
+      currentDirection := .Down,
+      takenSteps := .One
+    } :: a
+  else
+    a
+
+private def HeatLossMap.startPoints (heatLossMap : HeatLossMap) : List (LeanAStar.StartPoint (PathNode heatLossMap)) :=
+  heatLossMap.start.map λx ↦ {start := x, initial_costs := heatLossMap[x.coordinate]}
 
 def part1 (heatLossMap : HeatLossMap) : Option Nat :=
-  heatLossMap.findPath (OpenSet.start heatLossMap) Std.HashSet.empty
+  have : Add Nat := inferInstance
+  LeanAStar.findLowestCosts heatLossMap.startPoints
 
 ------------------------------------------------------------------------------------
 
@@ -421,4 +347,7 @@ private def testData := "2413432311323
 
 #eval match parse testData with
 | .error _ => none
-| .ok m => HeatLossMap.findPath (OpenSet.start m) Std.HashSet.empty
+| .ok m =>
+  have : Add Nat := inferInstance
+  let r : Option Nat := LeanAStar.findLowestCosts m.startPoints
+  r