horse/prob: basic probability distribution stuff

horse/prob/pmf.kk

@@ -0,0 +1,36 @@
+module horse/prob/pmf
+
+import std/core/list
+
+// Discrete-support probability distribution implemented as a list with the invariant
+// that support is always given in increasing order.
+pub type pmf<s, v>
+  Event(s: s, v: v, next: pmf<s, v>)
+  End
+
+// Add an independent event to the distribution.
+pub fun add(p: pmf<s, v>, s: s, v: v, ?s/cmp: (a: s, b: s) -> order, ?v/(+): (new: v, old: v) -> e v): e pmf<s, v>
+  match p
+    End -> Event(s, v, End)
+    Event(s', v', next) -> match s.cmp(s')
+      Lt -> Event(s, v, Event(s', v', next))
+      Eq -> Event(s, v + v', next)
+      Gt -> Event(s', v', add(next, s, v))
+
+// Replace an event in the distribution.
+pub inline fun set(p: pmf<s, v>, s: s, v: v, ?s/cmp: (a: s, b: s) -> order): e pmf<s, v>
+  p.add(s, v, cmp, fn(new, old) new)
+
+// Construct a pmf from a list of (support, value) entries.
+pub fun list/pmf(l: list<(s, v)>, ?s/cmp: (a: s, b: s) -> order, ?v/(+): (new: v, old: v) -> e v): e pmf<s, v>
+  l.foldl(End) fn(p, (s, v)) p.add(s, v)
+
+// Fold over the entries of the distribution.
+pub tail fun foldl(p: pmf<s, v>, init: a, f: (a, s, v) -> e a): e a
+  match p
+    End -> init
+    Event(s, v, next) -> foldl(next, f(init, s, v), f)
+
+// Convert the distribution to a list of entries.
+pub fun pmf/list(p: pmf<s, v>): list<(s, v)>
+  p.foldl(Nil) fn(l, s, v) Cons((s, v), l)