|
| 1 | +import Std |
| 2 | +open Std |
| 3 | + |
| 4 | +namespace LinkedList |
| 5 | + |
| 6 | +variable {α σ : Type} [BEq α] [Inhabited α] |
| 7 | + |
| 8 | +structure Node (α : Type) where |
| 9 | + prev : Option Nat |
| 10 | + value : α |
| 11 | + next : Option Nat |
| 12 | + deriving Inhabited |
| 13 | + |
| 14 | +structure LinkedList σ α where |
| 15 | + nodes : ST.Ref σ (Array (Node α)) |
| 16 | + free : ST.Ref σ (Array Nat) |
| 17 | + head : ST.Ref σ (Option Nat) |
| 18 | + tail : ST.Ref σ (Option Nat) |
| 19 | + length : ST.Ref σ Nat |
| 20 | + |
| 21 | +def LinkedList.empty : ST σ (LinkedList σ α) := do |
| 22 | + return { |
| 23 | + nodes := ← ST.mkRef #[], |
| 24 | + free := ← ST.mkRef #[], |
| 25 | + head := ← ST.mkRef none, |
| 26 | + tail := ← ST.mkRef none, |
| 27 | + length := ← ST.mkRef 0 |
| 28 | + } |
| 29 | + |
| 30 | +private def buildNode (value : α) (list : LinkedList σ α) : ST σ Nat := do |
| 31 | + let node : Node α := { prev := none, value := value, next := none} |
| 32 | + match (← list.free.get).back? with |
| 33 | + | none => |
| 34 | + let id := (← list.nodes.get).size |
| 35 | + list.nodes.modify (·.push node) |
| 36 | + return id |
| 37 | + | some id => |
| 38 | + list.nodes.modify (·.set! id node) |
| 39 | + list.free.modify (·.pop) |
| 40 | + return id |
| 41 | + |
| 42 | +private def LinkedList.unlinkPrev (node : Node α) (list : LinkedList σ α) : ST σ Unit := do |
| 43 | + match node.prev with |
| 44 | + | some p => |
| 45 | + list.nodes.modify (·.modify p (λ pn => { pn with next := node.next })) |
| 46 | + | none => |
| 47 | + list.head.set node.next |
| 48 | + |
| 49 | +private def LinkedList.unlinkNext (node : Node α) (list : LinkedList σ α) : ST σ Unit := do |
| 50 | + match node.next with |
| 51 | + | some n => |
| 52 | + list.nodes.modify (·.modify n (λ np => { np with prev := node.prev })) |
| 53 | + | none => |
| 54 | + list.tail.set node.prev |
| 55 | + |
| 56 | +private def LinkedList.unlink (node : Node α) (list : LinkedList σ α) : ST σ Unit := do |
| 57 | + list.unlinkPrev node |
| 58 | + list.unlinkNext node |
| 59 | + |
| 60 | +def LinkedList.count (list : LinkedList σ α) : ST σ Nat := |
| 61 | + list.length.get |
| 62 | + |
| 63 | +def LinkedList.push (value : α) (list : LinkedList σ α) : ST σ Unit := do |
| 64 | + let id ← buildNode value list |
| 65 | + match ← list.tail.get with |
| 66 | + | none => |
| 67 | + list.head.set (some id) |
| 68 | + list.tail.set (some id) |
| 69 | + list.length.modify (· + 1) |
| 70 | + | some t => |
| 71 | + list.nodes.modify (λ ns => |
| 72 | + ns.modify t (λ n => { n with next := some id }) |
| 73 | + |>.modify id (λ n => { n with prev := some t }) |
| 74 | + ) |
| 75 | + list.tail.set (some id) |
| 76 | + list.length.modify (· + 1) |
| 77 | + |
| 78 | +def LinkedList.unshift (value : α) (list : LinkedList σ α) : ST σ Unit := do |
| 79 | + let id ← buildNode value list |
| 80 | + match ← list.head.get with |
| 81 | + | none => |
| 82 | + list.head.set (some id) |
| 83 | + list.tail.set (some id) |
| 84 | + list.length.modify (· + 1) |
| 85 | + | some h => |
| 86 | + list.nodes.modify (λ ns => |
| 87 | + ns.modify h (λ n => { n with prev := some id }) |
| 88 | + |>.modify id (λ n => { n with next := some h }) |
| 89 | + ) |
| 90 | + list.head.set (some id) |
| 91 | + list.length.modify (· + 1) |
| 92 | + |
| 93 | +def LinkedList.pop (list : LinkedList σ α) : ST σ (Option α) := do |
| 94 | + match (← list.length.get) with |
| 95 | + | 0 => return none |
| 96 | + | n + 1 => |
| 97 | + let id := (← list.tail.get).get! |
| 98 | + let node := (← list.nodes.get)[id]! |
| 99 | + list.unlink node |
| 100 | + list.length.set n |
| 101 | + list.free.modify (·.push id) |
| 102 | + return node.value |
| 103 | + |
| 104 | +def LinkedList.shift (list : LinkedList σ α) : ST σ (Option α) := do |
| 105 | + match (← list.length.get) with |
| 106 | + | 0 => return none |
| 107 | + | n + 1 => |
| 108 | + let id := (← list.head.get).get! |
| 109 | + let node := (← list.nodes.get)[id]! |
| 110 | + list.unlink node |
| 111 | + list.length.set n |
| 112 | + list.free.modify (·.push id) |
| 113 | + return node.value |
| 114 | + |
| 115 | +def LinkedList.delete (value : α) (list : LinkedList σ α) : ST σ Unit := do |
| 116 | + let mut crtId ← list.head.get |
| 117 | + while crtId.isSome do |
| 118 | + let id := crtId.get! |
| 119 | + let node := (← list.nodes.get)[id]! |
| 120 | + if node.value == value then |
| 121 | + list.unlink node |
| 122 | + list.length.modify (· - 1) |
| 123 | + list.free.modify (·.push id) |
| 124 | + break |
| 125 | + crtId := node.next |
| 126 | + |
| 127 | +end LinkedList |
0 commit comments