dijkstra.arr

include string-dict

data ControlFlow<a, b>: next(v :: a) | brk(v :: b) end
fun loop<a, b>(func :: (a -> ControlFlow<a, b>), init :: a) -> b:
  cases(ControlFlow) func(init):
    | brk(v) => v
    | next(v) => loop(func, v)
  end
end

data LeftHeap<a>:
  | lheap-empty with:
    method rank(self) -> Number: 
      0 
    end
  | lheap-node(value :: { Number; Number; a; LeftHeap<a>; LeftHeap<a> }) with:
    method rank(self) -> Number:
      {r; _; _; _; _} = self.value
      r
    end
sharing:
  method glue(self, xp :: Number, x :: a, other :: LeftHeap<a>) -> LeftHeap<a>:
    ask:
      | self.rank() >= other.rank() then: lheap-node({other.rank() + 1; xp; x; self; other})
      | otherwise: lheap-node({self.rank() + 1; xp; x; other; self})
    end
  end,
  method merge(self, other :: LeftHeap<a>) -> LeftHeap<a>:
    ask: 
      | is-lheap-empty(self) then: other
      | is-lheap-empty(other) then: self
      | otherwise:
        {_; xp; x; a1; b1} = self.value
        {_; yp; y; a2; b2} = other.value
        if xp <= yp:
          a1.glue(xp, x, b1.merge(other))
        else:
          a2.glue(yp, y, self.merge(b2))
        end
    end
  end,
  method enqueue(self, v :: {Number; a}) -> LeftHeap<a>:
    lheap-node({ 1; v.{0}; v.{1}; lheap-empty; lheap-empty }).merge(self)
  end,
  method dequeue(self) -> Option<{Number; a; LeftHeap<a>}>:
    cases(LeftHeap) self:
      | lheap-empty => none
      | lheap-node(value) =>
        {_; xp; x; l; r} = value
        some({xp; x; l.merge(r)})
    end
  end
end

fun lheap-from-raw-array<a>(arr :: RawArray<{Number; a}>) -> LeftHeap<a>:
  al = raw-array-length(arr)
  for loop({hp; i} from {lheap-empty; 0}):
    ask:
      | i >= al then: brk(hp)
      | otherwise: next({hp.enqueue(raw-array-get(arr, i)); i + 1})
    end
  end
end

lheap = {
  make0: lam():
    lheap-empty
  end,
  make1: lam(t0):
    lheap-empty.enqueue(t0)
  end,
  make2: lam(t0, t1):
    lheap-empty.enqueue(t0).enqueue(t1)
  end,
  make3: lam(t0, t1, t2):
    lheap-empty.enqueue(t0).enqueue(t1).enqueue(t2)
  end,
  make4: lam(t0, t1, t2, t3):
    lheap-empty.enqueue(t0).enqueue(t1).enqueue(t2).enqueue(t3)
  end,
  make5: lam(t0, t1, t2, t3, t4):
    lheap-empty.enqueue(t0).enqueue(t1).enqueue(t2).enqueue(t3).enqueue(t4)
  end,
  make: lheap-from-raw-array
}

data Joint: joint(destination-node :: String, distance :: Number) end
data Graph: 
  graph(nodes :: StringDict<List<Joint>>) with:
    method find-path(self, a :: String, b :: String) -> Option<{ Number; List<Option<String>> }>:
      fun get-path-from-dict(start-node, dirs-dct):
        for loop(l from [list: start-node]):
          cases(Option) dirs-dct.get-value(l.get(0)).{1}:
            | none => brk(l)
            | some(n) => next(link(n, l))
          end
        end
      end
  
      fun enqueue-next(p-queue, node):
        for fold(q from p-queue, jnt from self.nodes.get-value(node)):
          q.enqueue({ jnt.distance; { jnt.destination-node; some(node) } })
        end
      end
  
      ask:
        | a == b then: some({ 0; [list: a] } )
        | not(self.nodes.has-key(a)) or not(self.nodes.has-key(b)) then: none
        | otherwise:
          for loop({p-queue; dirs-dct} from { [lheap: { 0; {a; none} }]; [string-dict:] }):
            cases(Option) p-queue.dequeue():
              | none => brk(none)
              | some(res) =>
                { joint-dist; {node; src }; shadow p-queue } = res
                
                dist = cases(Option) src:
                  | none => 0
                  | some(it) => dirs-dct.get-value(it).{0}
                end + joint-dist
                
                if dirs-dct.get(node).and-then({(it): it.{0} <= dist}).or-else(false): 
                  next({p-queue; dirs-dct})
                else if node == b:
                  brk(some({ dist; get-path-from-dict(node, dirs-dct.set(node, { dist; src })) }))
                else:
                  next({ enqueue-next(p-queue, node); dirs-dct.set(node, { dist; src }) })
                end
            end
          end
      end
    end
end

check:
  g = graph(
    [string-dict:
      "A", [list: joint("B", 10)],
      "B", [list: joint("A", 10), joint("C", 6), joint("E", 8)],
      "C", [list: joint("B", 6), joint("D", 28), joint("E", 6)],
      "D", [list: joint("C", 28), joint("E", 14), joint("G", 10)],
      "E", [list: joint("B", 8), joint("C", 6), joint("D", 14), joint("F", 13)],
      "F", [list: joint("E", 13), joint("G", 15)],
      "G", [list: joint("F", 15), joint("D", 10)],
      "H", [list: joint("I", 35)],
      "I", [list: joint("H", 35)]
    ]
  )
  g.find-path("A", "A") is some({ 0; [list: "A"] })
  g.find-path("A", "B") is some({ 10; [list: "A", "B"] })
  g.find-path("A", "C") is some({ 16; [list: "A", "B", "C"] })
  g.find-path("A", "G") is some({ 42; [list: "A", "B", "E", "D", "G"] })
  g.find-path("H", "I") is some({ 35; [list: "H", "I"] })
  g.find-path("I", "H") is some({ 35; [list: "I", "H"] })
  g.find-path("A", "H") is none
  g.find-path("I", "A") is none
end