microkanren with semirings

kanrig.rkt

#lang racket

(struct rig (zero? sum product) #:prefab)
(define the-rig (make-parameter #f))
(define (the-rig-zero? x) ((rig-zero? (the-rig)) x))
(define (the-rig-sum args) ((rig-sum (the-rig)) args))
(define (the-rig-product args) ((rig-product (the-rig)) args))
(define (the-rig-+ . args) (the-rig-sum args))
(define (the-rig-* . args) (the-rig-product args))
(define (the-rig-zero) (the-rig-+))
(define (the-rig-one) (the-rig-*))


;; Racket's built-in lazy streams are synchronous, or orderly: exactly one delay between
;; successive elements.
;;
;; However, logic programming search may need to combine two parallel streams which do not
;; generate elements in lockstep, at the same rate; indeed, one stream may be
;; unproductive, and never generate elements at all, but this should not halt exploration
;; of the other stream.
;;
;; Instead we use asynchronous, or disorderly, streams. A disorderly stream is either:
;;
;; - empty, ()
;; - a pair (x . xs) of an element x and a stream xs.
;; - a promise that produces a stream.
;;
;; Note that a stream may begin with *arbitrarily many* delays, ie. arbitrarily nested
;; promises. Each delay gives us the opportunity to go explore somewhere else before
;; proceeding further, thus avoiding unproductive livelock.
;;
;; NOTE TO FUTURE SELF: Disorderly streams can be implemented as orderly stream of
;; *vectors* of elements:
;;
;; - a pair (x . xs) becomes (stream-cons (vector x) xs)
;; - a delayed stream (delay xs) becomes (stream-cons (vector) xs)
;;
;; Is this useful? Speculatively, it might:
;;
;; a) take advantage of locality by lumping groups of elements together; or
;;
;; b) allow a more databasey approach where we see these vectors as deltas to tables and
;;    build indices, semi-naive evaluation style. But how does this interact with the
;;    sideways information passing style with which miniKanren implements conj?

(define stream? (or/c null? pair? promise?))

(define (stream . xs) xs)
(define empty-stream '())
(define-syntax-rule (stream-delay e) (delay e))

(define/contract (stream-expose s)
  (-> stream? (or/c null? (cons/c any/c stream?)))
  (if (promise? s) (stream-expose (force s)) s))

(define/contract (stream->list s)
  (-> stream? list?)
  (match (stream-expose s)
    ['() '()]
    [(cons x xs) (cons x (stream->list xs))]))

(define/contract (stream-plus s1 s2)
  (-> stream? stream? stream?)
  (match s1
    ['() s2]
    [(cons x xs) (cons x (stream-plus s2 xs))]
    ;; take the hint; try s2 first.
    [(? promise?) (delay (stream-plus s2 (force s1)))]))

(define/contract (stream-interleave streams)
  (-> (listof stream?) stream?)
  (foldr stream-plus '() streams))

(define/contract (stream-bind s cont)
  (-> stream? (-> any/c stream?) stream?)
  (match s
    ['() '()]
    [(cons x xs) (stream-plus (cont x) (stream-bind xs cont))]
    [(? promise?) (delay (stream-bind (force s) cont))]))


;; A LOGIC PROGRAM is represented as a function:
;;
;;    f: factor subst ↦ results
;;
;; subst is a map from vars to terms
;; factor is a semiring value used to multiply all results by
;; results is a stream of results, which contain subst/semiring-value pairs.


;; TERMs can be anything, but are typically made of conses, symbols, and vars ;;
(define term? any/c)


;; UNIFICATION VARs are symbols beginning with ? ;;
(define (var? x) (and (symbol? x) (string-prefix? (symbol->string x) "?")))
(define/contract (var x) (-> symbol? var?)
  (string->symbol (format "?~a" x)))
(define/contract (var-name x) (-> var? symbol?)
  (string->symbol (string-replace (symbol->string x) "?" "" #:all? #f)))


;; RESULTs are pairs (value . subst) ;;
(define-match-expander result
  (syntax-rules () [(_ value subst) (cons value subst)])
  (syntax-rules () [(_ value subst) (cons value subst)]))
(define result-value car)
(define result-subst cdr)


;; SUBSTITUTIONs are hashes from variable names to terms ;;

;; (define subst? (hash/c var? term? #:immutable #t #:flat? #t))
;; (define/contract subst-empty subst? (hash))
;; (define (subst-contains? s x) (hash-has-key? s x))
;; (define (subst-ref s x) (hash-ref s x))
;; (define/contract (subst-extend s x t)
;;   (-> subst? var? term? subst?)
;;   (hash-set s x t))

(define subst? list?)
(define subst-empty '())
(define (subst-contains? s x) (pair? (assoc x s)))
(define (subst-ref s x) (caddr (assoc x s)))
(define (subst-extend s x t) (cons (list x '= t) s))


;;;; KANRIG PROGRAMS ;;;;
;;
;; semiring terms V             TODO implement this via syntax?
;;   (zero)
;;   (one)
;;   (plus ...)
;;   (times ...)
;;
;; terms M,N
;;
;; programs P,Q
;;   (== M N)
;;   (yield V)
;;   (disj P ...)
;;   (conj P ...)
;;   (exists (x ...) P ...)         binds fresh x... and products P...
;;   (conde (P ...) ...)            sum of products
;;   (delaye P)                     delays evaluation - necessary?

(define ((disj . Ps) factor subst)
  (stream-interleave (for/list ([P Ps]) (P factor subst))))

;; yield always wants to be eta-delayed so that V can have access to the rig.
(define-syntax-rule (yield V)
  (lambda (factor subst)
    (stream (result (the-rig-* factor V) subst))))

(define (conj . Ps) (foldr mplus (yield (the-rig-one)) Ps))
(define ((mplus P1 P2) factor subst)
  (stream-bind
   (P1 factor subst)
   (match-lambda [(result factor subst) (P2 factor subst)])))

(define-syntax exists
  (syntax-rules ()
    [(_ () body ...) (conj body ...)]
    [(_ (x y ...) body ...)
     (call/fresh (lambda (x) (exists (y ...) body ...)))]))
(define ((call/fresh f) factor subst)
  (define x (var (gensym)))
  ((f x) factor subst))

(define-syntax-rule (conde (P ...) ...)
  (disj (conj P ...) ...))

;; Necessary for recursive definitions.
(define ((delaye P) factor subst)
  (stream-delay (P factor subst)))


;;;; UNIFICATION ;;;;
(define ((== t u) factor subst)
  (match (unify t u subst)
    [#f empty-stream]
    [subst (stream (result factor subst))]))

;; NB. no occurs check. Add one?
(define (unify t u s)
  (match* ((walk s t) (walk s u))
    [(t u) #:when (equal? t u) s]
    [((? var? x) t) (subst-extend s x t)]
    [(t (? var? x)) (subst-extend s x t)]
    [((cons t1 t2) (cons u1 u2))
     (define s1 (unify t1 u1 s))
     (and s1 (unify t2 u2 s1))]
    [(_ _) #f]))

(define (walk s t)
  (if (and (var? t) (subst-contains? s t))
      (walk s (subst-ref s t))
      t))


;;;; EVALUATION & TESTS ;;;;

(define (results P) (P (the-rig-one) subst-empty))
(define (all-results P) (stream->list (results P)))

(define (pretty P #:n [n #f])
  (let loop ([i 0]
             [rs (results P)])
    (if (eqv? i n) (displayln "...")
        (match (stream-expose rs)
          ['() (void)]
          [(cons r rs)
           (displayln (reverse r))      ; TODO: abstraction breaking
           (loop (+ i 1) rs)]))))

;; EXAMPLES ;;
(define minplus-rig
  (rig (lambda (x) (= +inf.0 x))
       (lambda (xs) (apply min xs))
       (lambda (xs) (apply + xs))))

(the-rig minplus-rig)

(define foo (var (gensym 'foo)))

;;    3     7
;; a --- b --- c
;;  \         /
;;   ---------
;;       8
(define (edgeo x y)
  (conde
   [(== x 'a) (== y 'b) (yield 3)]
   [(== x 'b) (== y 'c) (yield 7)]
   [(== x 'a) (== y 'c) (yield 8)]))

;; Why does this terminate, even without delaye? Because the graph is finite and acyclic
;; (and thus "forward-finite": for any vertex v, there are finitely many paths out of it)
;; and it always looks for an edge first; so if it gets to a node with no out-edges,
;; (edgeo x z) fails and (patho z y) never runs.
(define (patho x y)
  (delaye ;; technically not necessary for the reasons given above
   (disj (edgeo x y)
         (exists (z) (edgeo x z) (patho z y)))))

;;; This, on the other hand, needs delaye to be productive; even *with* delaye, it doesn't
;;; terminate after finding all solutions.
(define (bad-patho x y)
  (delaye
   (disj (edgeo x y)
         (exists (z) (bad-patho x z) (edgeo z y)))))