# HG changeset patch
# User Steve Losh <steve@stevelosh.com>
# Date 1461170018 0
# Node ID 1ab41e0128dc4a3bc94e905bc10deb6e97a6533d
# Parent  67535b9c3b86ef02079af9c24f140b4caf897e27
Add the TRY* instructions to compilation

Still need to implement the actual bytecode.

diff -r 67535b9c3b86 -r 1ab41e0128dc src/wam/bytecode.lisp
--- a/src/wam/bytecode.lisp	Tue Apr 19 14:00:32 2016 +0000
+++ b/src/wam/bytecode.lisp	Wed Apr 20 16:33:38 2016 +0000
@@ -35,7 +35,10 @@
     (+opcode-proceed+ 1)
     (+opcode-allocate+ 2)
     (+opcode-deallocate+ 1)
-    (+opcode-done+ 1)))
+    (+opcode-done+ 1)
+    (+opcode-try+ 2)
+    (+opcode-retry+ 2)
+    (+opcode-trust+ 1)))
 
 
 (defun* opcode-name ((opcode opcode))
@@ -66,7 +69,10 @@
     (+opcode-proceed+ "PROCEED")
     (+opcode-allocate+ "ALLOCATE")
     (+opcode-deallocate+ "DEALLOCATE")
-    (+opcode-done+ "DONE")))
+    (+opcode-done+ "DONE")
+    (+opcode-try+ "TRY")
+    (+opcode-retry+ "RETRY")
+    (+opcode-trust+ "TRUST")))
 
 (defun* opcode-short-name ((opcode opcode))
   (:returns string)
@@ -97,5 +103,8 @@
     (+opcode-proceed+ "PROC")
     (+opcode-allocate+ "ALOC")
     (+opcode-deallocate+ "DEAL")
-    (+opcode-done+ "DONE")))
+    (+opcode-done+ "DONE")
+    (+opcode-try+ "TRYM")
+    (+opcode-retry+ "RTRY")
+    (+opcode-trust+ "TRST")))
 
diff -r 67535b9c3b86 -r 1ab41e0128dc src/wam/compiler.lisp
--- a/src/wam/compiler.lisp	Tue Apr 19 14:00:32 2016 +0000
+++ b/src/wam/compiler.lisp	Wed Apr 20 16:33:38 2016 +0000
@@ -474,12 +474,12 @@
 
 (defun tokenize-program-term
     (term permanent-variables reserved-variables reserved-arity)
-  "Tokenize `term` as a program term, returning its tokens, functor, and arity."
-  (tokenize-term term
-                 permanent-variables
-                 reserved-variables
-                 reserved-arity
-                 #'flatten-program))
+  "Tokenize `term` as a program term, returning its tokens."
+  (values (tokenize-term term
+                         permanent-variables
+                         reserved-variables
+                         reserved-arity
+                         #'flatten-program)))
 
 (defun tokenize-query-term
     (term permanent-variables &optional reserved-variables reserved-arity)
@@ -604,7 +604,7 @@
       (handle-stream body-tokens))))
 
 
-;;;; UI
+;;;; Compilation
 (defun find-variables (terms)
   "Return the set of variables in `terms`."
   (remove-duplicates (tree-collect #'variable-p terms)))
@@ -640,20 +640,6 @@
       (find-shared-variables (list head body-first)))))
 
 
-(defun mark-label (wam functor arity store)
-  "Set the code label `(functor . arity)` to point at the next space in `store`."
-  ;; todo make this less ugly
-  (setf (wam-code-label wam (wam-ensure-functor-index wam (cons functor arity)))
-        (fill-pointer store)))
-
-
-(defun make-query-code-store ()
-  (make-array 64
-    :fill-pointer 0
-    :adjustable t
-    :element-type 'code-word))
-
-
 (defun compile-clause (wam store head body)
   "Compile the clause directly into `store` and return the permanent variables.
 
@@ -677,13 +663,10 @@
                 (1- (length (car body)))))
          (head-tokens
            (when head
-             (multiple-value-bind (tokens functor arity)
-                 (tokenize-program-term head
-                                        permanent-variables
-                                        head-variables
-                                        head-arity)
-               (mark-label wam functor arity store) ; TODO: this is ugly
-               tokens)))
+             (tokenize-program-term head
+                                    permanent-variables
+                                    head-variables
+                                    head-arity)))
          (body-tokens
            (when body
              (append
@@ -717,6 +700,14 @@
          (code-push-instruction! store +opcode-done+))))
     permanent-variables))
 
+
+;;; Queries
+(defun make-query-code-store ()
+  (make-array 64
+    :fill-pointer 0
+    :adjustable t
+    :element-type 'code-word))
+
 (defun compile-query (wam query)
   "Compile `query` into a fresh array of bytecode.
 
@@ -729,13 +720,65 @@
          (permanent-variables (compile-clause wam store nil query)))
     (values store permanent-variables)))
 
-(defun compile-program (wam rule)
-  "Compile `rule` into the WAM's code store.
+
+;;; Rules
+(defun mark-label (wam functor arity address)
+  "Set the code label `functor`/`arity` to point at `address`."
+  (setf (wam-code-label wam functor arity) address))
+
+(defun find-arity (rule)
+  (let ((head (first rule)))
+    (cond
+      ((null head) (error "Rule ~S has a NIL head." rule))
+      ((atom head) 0) ; constants are 0-arity
+      (t (1- (length head))))))
 
-  `rule` should be a clause consisting of a head term and zero or more body
-  terms.  A rule with no body is called a fact.
+(defun check-rules (rules)
+  (let* ((predicates (mapcar #'caar rules))
+         (arities (mapcar #'find-arity rules))
+         (functors (zip predicates arities)))
+    (assert (= 1 (length (remove-duplicates functors :test #'equal))) ()
+      "Must add exactly 1 predicate at a time (got: ~S)."
+      functors)
+    (values (first predicates) (first arities))))
+
+(defun compile-rules (wam rules)
+  "Compile `rules` into the WAM's code store.
+
+  Each rule in `rules` should be a clause consisting of a head term and zero or
+  more body terms.  A rule with no body is called a fact.
 
   "
-  (compile-clause wam (wam-code wam) (first rule) (rest rule))
+  (assert rules () "Cannot compile an empty program.")
+  (*let ((code (wam-code wam))
+         (previous-jump nil)
+         ((:values functor arity) (check-rules rules)))
+    (labels
+        ((fill-jump (address)
+           (when previous-jump
+             (setf (aref code (1+ previous-jump)) address))
+           (setf previous-jump address))
+         (push-branch-instruction (first-p last-p)
+           (cond
+             (first-p
+              (fill-jump (code-push-instruction! code +opcode-try+ 999)))
+             (last-p
+              (fill-jump (code-push-instruction! code +opcode-trust+)))
+             (t
+              (fill-jump (code-push-instruction! code +opcode-retry+ 999))))))
+      ;; Mark the label to point at where we're about to stick the code.
+      ;; TODO: this is ugly
+      (mark-label wam functor arity (fill-pointer code))
+      (if (= 1 (length rules))
+        ;; Single-clause rules don't need to bother setting up a choice point.
+        (destructuring-bind ((head . body)) rules
+          (compile-clause wam code head body))
+        ;; Otherwise we need to loop through each of the clauses, pushing their
+        ;; choice point instruction first, then their actual code.
+        (loop :for ((head . body) . remaining) :on rules
+              :for first-p = t :then nil
+              :do
+              (push-branch-instruction first-p (null remaining))
+              (compile-clause wam code head body)))))
   (values))
 
diff -r 67535b9c3b86 -r 1ab41e0128dc src/wam/constants.lisp
--- a/src/wam/constants.lisp	Tue Apr 19 14:00:32 2016 +0000
+++ b/src/wam/constants.lisp	Wed Apr 20 16:33:38 2016 +0000
@@ -48,10 +48,10 @@
   :documentation "Bitmask for the functor arity bits.")
 
 
-(define-constant +register-count+ 16
+(define-constant +register-count+ 2048
   :documentation "The number of registers the WAM has available.")
 
-(define-constant +maximum-arity+ (1- (expt 2 +functor-arity-width+))
+(define-constant +maximum-arity+ 1024
   :documentation "The maximum allowed arity of functors.")
 
 
@@ -60,13 +60,6 @@
   "The maximum size (in bytes of bytecode) a query may compile to.")
 
 
-(define-constant +tag-local-register+ #b0
-  :documentation "A local register (X_n or A_n).")
-
-(define-constant +tag-stack-register+ #b1
-  :documentation "A stack register (Y_n).")
-
-
 (define-constant +stack-word-size+ 16
   :documentation "Size (in bits) of each word in WAM stack.")
 
@@ -79,10 +72,16 @@
   ;; too large.
   :documentation "Maximum size of the WAM stack.")
 
-(define-constant +stack-frame-size-limit+ (+ 3 +register-count+)
+(define-constant +stack-frame-size-limit+ (+ 7 +register-count+)
   :documentation "The maximum size, in stack frame words, that a stack frame could be.")
 
 
+(define-constant +trail-limit+ (expt 2 +stack-word-size+)
+  ;; The trail's fill pointer is stored inside choice frames on the stack, so it
+  ;; needs to be able to fit inside a stack word.
+  :documentation "The maximum number of variables that may exist in the trail.")
+
+
 ;;;; Opcodes
 ;;; Program
 (define-constant +opcode-noop+ 0)
@@ -115,6 +114,9 @@
 (define-constant +opcode-allocate+ 21)
 (define-constant +opcode-deallocate+ 22)
 (define-constant +opcode-done+ 23)
+(define-constant +opcode-try+ 24)
+(define-constant +opcode-retry+ 25)
+(define-constant +opcode-trust+ 26)
 
 
 ;;;; Debug Config
diff -r 67535b9c3b86 -r 1ab41e0128dc src/wam/dump.lisp
--- a/src/wam/dump.lisp	Tue Apr 19 14:00:32 2016 +0000
+++ b/src/wam/dump.lisp	Wed Apr 20 16:33:38 2016 +0000
@@ -197,7 +197,6 @@
           (second arguments)
           (first arguments)))
 
-
 (defmethod instruction-details ((opcode (eql +opcode-call+)) arguments functor-list)
   (format nil "CALL~A      ; ~A"
           (pretty-arguments arguments)
@@ -244,7 +243,7 @@
 
 (defun dump-wam-registers (wam)
   (format t "REGISTERS:~%")
-  (format t  "~5@A ->~6@A~%" "S" (wam-s wam))
+  (format t  "~5@A ->~6@A~%" "S" (wam-subterm wam))
   (loop :for i :from 0
         :for reg :across (wam-local-registers wam)
         :for contents = (when (not (= reg (1- +heap-limit+)))
@@ -286,6 +285,10 @@
   (dump-labels wam)
   (dump-code wam))
 
+(defun dump-wam-code (wam)
+  (with-slots (code) wam
+    (dump-code-store wam code +maximum-query-size+ (length code))))
+
 (defun dump-wam-full (wam)
   (dump-wam wam 0 (length (wam-heap wam)) -1))
 
diff -r 67535b9c3b86 -r 1ab41e0128dc src/wam/interpreter.lisp
--- a/src/wam/interpreter.lisp	Tue Apr 19 14:00:32 2016 +0000
+++ b/src/wam/interpreter.lisp	Wed Apr 20 16:33:38 2016 +0000
@@ -249,7 +249,7 @@
 (define-instruction %get-structure-local ((wam wam)
                                           (functor functor-index)
                                           (register register-index))
-  (with-accessors ((mode wam-mode) (s wam-s)) wam
+  (with-accessors ((mode wam-mode) (s wam-subterm)) wam
     (let* ((addr (deref wam (wam-local-register wam register)))
            (cell (wam-heap-cell wam addr)))
       (cond
@@ -305,11 +305,11 @@
      (register register-index))
   (ecase (wam-mode wam)
     (:read (setf (%wam-register% wam register)
-                 (wam-s wam)))
+                 (wam-subterm wam)))
     (:write (->> (push-unbound-reference! wam)
               (nth-value 1)
               (setf (%wam-register% wam register)))))
-  (incf (wam-s wam)))
+  (incf (wam-subterm wam)))
 
 (define-instructions (%unify-value-local %unify-value-stack)
     ((wam wam)
@@ -317,12 +317,12 @@
   (ecase (wam-mode wam)
     (:read (unify! wam
                    (%wam-register% wam register)
-                   (wam-s wam)))
+                   (wam-subterm wam)))
     (:write (wam-heap-push! wam
                             (->> register
                               (%wam-register% wam)
                               (wam-heap-cell wam)))))
-  (incf (wam-s wam)))
+  (incf (wam-subterm wam)))
 
 (define-instructions (%get-variable-local %get-variable-stack)
     ((wam wam)
@@ -369,7 +369,7 @@
 (define-instruction %deallocate ((wam wam))
   (setf (wam-program-counter wam)
         (wam-stack-frame-cp wam))
-  (wam-stack-pop-environment! wam))
+  (wam-stack-pop-frame! wam))
 
 
 ;;;; Running
diff -r 67535b9c3b86 -r 1ab41e0128dc src/wam/types.lisp
--- a/src/wam/types.lisp	Tue Apr 19 14:00:32 2016 +0000
+++ b/src/wam/types.lisp	Wed Apr 20 16:33:38 2016 +0000
@@ -16,6 +16,9 @@
 (deftype stack-index ()
   `(integer 0 ,(1- +stack-limit+)))
 
+(deftype trail-index ()
+  `(integer 0 ,(1- +trail-limit+)))
+
 (deftype register-index ()
   `(integer 0 ,(1- +register-count+)))
 
@@ -39,14 +42,17 @@
 
 
 (deftype opcode ()
-  '(integer 0 23))
+  '(integer 0 26))
 
 
 (deftype stack-frame-size ()
   `(integer 3 ,+stack-frame-size-limit+))
 
+(deftype stack-choice-size ()
+  `(integer 7 ,+stack-frame-size-limit+))
+
 (deftype stack-frame-argcount ()
-  `(integer 0 ,+register-count+))
+  'arity)
 
 (deftype continuation-pointer ()
   'code-index)
@@ -54,9 +60,27 @@
 (deftype environment-pointer ()
   'stack-index)
 
-(deftype stack-word ()
+(deftype backtrack-pointer ()
+  'stack-index)
+
+
+(deftype stack-frame-word ()
   '(or
     environment-pointer ; CE
     continuation-pointer ; CP
     stack-frame-argcount ; N
-    heap-index)) ; YN
+    heap-index)) ; Yn
+
+(deftype stack-choice-word ()
+  '(or
+    environment-pointer ; CE
+    backtrack-pointer ; B
+    continuation-pointer ; CP, BP
+    stack-frame-argcount ; N
+    trail-index ; TR
+    heap-index))  ; An, H
+
+(deftype stack-word ()
+  '(or stack-frame-word stack-choice-word))
+
+
diff -r 67535b9c3b86 -r 1ab41e0128dc src/wam/ui.lisp
--- a/src/wam/ui.lisp	Tue Apr 19 14:00:32 2016 +0000
+++ b/src/wam/ui.lisp	Wed Apr 20 16:33:38 2016 +0000
@@ -8,15 +8,19 @@
      ,@body))
 
 
-(defun add-rule (rule)
-  (compile-program *database* rule))
+
+(defun add-rules (rules)
+  (compile-rules *database* rules))
 
 (defun perform-query (query step)
   (run-query *database* query step))
 
 
 (defmacro rule (&body body)
-  `(add-rule ',body))
+  `(add-rules '(,body)))
+
+(defmacro rules (&body rules)
+  `(add-rules ',rules))
 
 (defmacro query (&body body)
   `(perform-query ',body nil))
@@ -24,5 +28,8 @@
 (defmacro query-step (&body body)
   `(perform-query ',body t))
 
-(defun dump ()
-  (dump-wam-full *database*))
+
+(defun dump (&optional full-code)
+  (dump-wam-full *database*)
+  (when full-code
+    (dump-wam-code *database*)))
diff -r 67535b9c3b86 -r 1ab41e0128dc src/wam/wam.lisp
--- a/src/wam/wam.lisp	Tue Apr 19 14:00:32 2016 +0000
+++ b/src/wam/wam.lisp	Wed Apr 20 16:33:38 2016 +0000
@@ -63,26 +63,48 @@
                  :adjustable t
                  :element-type 'heap-index)
      :documentation "The unification stack.")
-   (s
-     :accessor wam-s
+   (trail
+     :reader wam-trail
+     :initform (make-array 64
+                 :fill-pointer 0
+                 :adjustable t
+                 :element-type 'heap-index)
+     :documentation "The trail of variables to unbind on backtracking.")
+   (number-of-arguments
+     :accessor wam-nargs
+     :initform 0
+     :type arity
+     :documentation "The Number of Arguments register (global var).")
+   (subterm
+     :accessor wam-subterm
      :initform nil
      :type (or null heap-index)
-     :documentation "The S register (address of next subterm to match).")
+     :documentation "The Subterm register (S).")
    (program-counter
      :accessor wam-program-counter
      :initform 0
      :type code-index
-     :documentation "The Program Counter into the WAM code store.")
+     :documentation "The Program Counter (P) into the WAM code store.")
    (continuation-pointer
      :accessor wam-continuation-pointer
      :initform 0
      :type code-index
-     :documentation "The Continuation Pointer into the WAM code store.")
+     :documentation "The Continuation Pointer (CP) into the WAM code store.")
    (environment-pointer
      :accessor wam-environment-pointer
      :initform 0
-     :type stack-index
-     :documentation "The Environment Pointer into the WAM stack.")
+     :type environment-pointer
+     :documentation "The Environment Pointer (E) into the WAM stack.")
+   (backtrack-pointer
+     :accessor wam-backtrack-pointer
+     :initform 0
+     :type backtrack-pointer
+     :documentation "The Backtrack Pointer (B) into the WAM stack.")
+   (heap-backtrack-pointer
+     :accessor wam-heap-backtrack-pointer
+     :initform 0
+     :type heap-index
+     :documentation "The Heap Backtrack Pointer (HB) into the WAM heap.")
    (mode
      :accessor wam-mode
      :initform nil
@@ -122,7 +144,74 @@
   (setf (aref (wam-heap wam) address) new-value))
 
 
+;;;; Trail
+(defun* wam-trail-pointer ((wam wam))
+  (:returns trail-index)
+  "Return the current trail pointer of the WAM."
+  (fill-pointer (wam-trail wam)))
+
+(defun* wam-trail-push! ((wam wam) (address heap-index))
+  (:returns (values heap-index trail-index))
+  "Push `address` onto the trail.
+
+  Returns the address and the trail address it was pushed to.
+
+  "
+  (with-slots (trail) wam
+    (if (= +trail-limit+ (fill-pointer trail))
+      (error "WAM trail exhausted.")
+      (values address (vector-push-extend address trail)))))
+
+(defun* wam-trail-pop! ((wam wam))
+  (:returns heap-index)
+  "Pop the top address off the trail and return it."
+  (vector-pop (wam-trail wam)))
+
+
 ;;;; Stack
+(defun* wam-stack-pointer ((wam wam))
+  (:returns stack-index)
+  "Return the current stack pointer of the WAM."
+  (fill-pointer (wam-stack wam)))
+
+
+(defun* wam-stack-word ((wam wam) (address stack-index))
+  (:returns stack-word)
+  "Return the stack word at the given address."
+  (aref (wam-stack wam) address))
+
+(defun (setf wam-stack-word) (new-value wam address)
+  (setf (aref (wam-stack wam) address) new-value))
+
+
+(defun* wam-stack-push! ((wam wam) (word stack-word))
+  (:returns (values stack-word stack-index))
+  "Push the word onto the WAM stack and increment the stack pointer.
+
+  Returns the word and the address it was pushed to.
+
+  "
+  (with-slots (stack) wam
+    (if (= +stack-limit+ (fill-pointer stack))
+      (error "WAM stack exhausted.")
+      (values word (vector-push-extend word stack)))))
+
+(defun* wam-stack-extend! ((wam wam) (words integer))
+  (:returns :void)
+  "Extend the WAM stack by the given number of words.
+
+  Each word is initialized to 0.
+
+  "
+  ;; TODO: this sucks, fix it
+  (with-slots (stack) wam
+    (repeat words
+      (if (= +stack-limit+ (fill-pointer stack))
+        (error "WAM stack exhausted.")
+        (vector-push-extend 0 stack))))
+  (values))
+
+
 ;;; Stack frames are laid out like so:
 ;;;
 ;;;     |PREV|
@@ -131,24 +220,8 @@
 ;;;     | N  |
 ;;;     | Y0 |
 ;;;     | .. |
-;;;     | YN |
+;;;     | Yn |
 ;;;     |NEXT| <-- fill-pointer
-
-(defun* wam-stack-pointer ((wam wam))
-  (:returns stack-index)
-  "Return the current stack pointer of the WAM."
-  (fill-pointer (wam-stack wam)))
-
-
-(defun* wam-stack-word ((wam wam) (address stack-index))
-  (:returns stack-index)
-  "Return the stack word at the given address."
-  (aref (wam-stack wam) address))
-
-(defun (setf wam-stack-word) (new-value wam address)
-  (setf (aref (wam-stack wam) address) new-value))
-
-
 (defun* wam-stack-frame-ce
     ((wam wam)
      &optional
@@ -173,6 +246,7 @@
   (:returns stack-frame-argcount)
   (wam-stack-word wam (+ 2 e)))
 
+
 (defun* wam-stack-frame-arg
     ((wam wam)
      (n register-index)
@@ -206,39 +280,128 @@
   (+ (wam-stack-frame-n wam e) 3))
 
 
-(defun* wam-stack-push! ((wam wam) (word stack-word))
-  (:returns (values stack-word stack-index))
-  "Push the word onto the WAM stack and increment the stack pointer.
+(defun* wam-stack-pop-frame! ((wam wam))
+  "Pop an environment (stack frame) off the WAM stack."
+  (let ((size (wam-stack-frame-size wam)))
+    (with-slots (stack environment-pointer) wam
+      (setf environment-pointer
+            (wam-stack-frame-ce wam environment-pointer)) ; E <- CE
+      (decf (fill-pointer stack) size)))) ; its fine
+
 
-  Returns the word and the address it was pushed to.
+;;; Choice point frames are laid out like so:
+;;;
+;;;         |PREV|
+;;;       0 | N  | <-- backtrack-pointer
+;;;       1 | CE |
+;;;       2 | CP | This is a bit different than the book.  We stick the
+;;;       3 | CB | arguments at the end of the frame instead of the beginning,
+;;;       4 | BP | so it's easier to retrieve the other values.
+;;;       5 | TR |
+;;;       6 | H  |
+;;;       7 | A0 |
+;;;         | .. |
+;;;     7+n | An |
+;;;         |NEXT| <-- fill-pointer
 
-  "
-  (with-slots (stack) wam
-    (if (= +stack-limit+ (fill-pointer stack))
-      (error "WAM stack exhausted.")
-      (values word (vector-push-extend word stack)))))
+(defun* wam-stack-choice-n
+    ((wam wam)
+     &optional
+     ((b backtrack-pointer)
+      (wam-backtrack-pointer wam)))
+  (:returns arity)
+  (wam-stack-word wam b))
+
+(defun* wam-stack-choice-ce
+    ((wam wam)
+     &optional
+     ((b backtrack-pointer)
+      (wam-backtrack-pointer wam)))
+  (:returns environment-pointer)
+  (wam-stack-word wam (+ b 1)))
 
-(defun* wam-stack-extend! ((wam wam) (words integer))
-  (:returns :void)
-  "Extend the WAM stack by the given number of words.
+(defun* wam-stack-choice-cp
+    ((wam wam)
+     &optional
+     ((b backtrack-pointer)
+      (wam-backtrack-pointer wam)))
+  (:returns continuation-pointer)
+  (wam-stack-word wam (+ b 2)))
+
+(defun* wam-stack-choice-cb
+    ((wam wam)
+     &optional
+     ((b backtrack-pointer)
+      (wam-backtrack-pointer wam)))
+  (:returns backtrack-pointer)
+  (wam-stack-word wam (+ b 3)))
 
-  Each word is initialized to 0.
+(defun* wam-stack-choice-bp
+    ((wam wam)
+     &optional
+     ((b backtrack-pointer)
+      (wam-backtrack-pointer wam)))
+  (:returns continuation-pointer)
+  (wam-stack-word wam (+ b 4)))
+
+(defun* wam-stack-choice-tr
+    ((wam wam)
+     &optional
+     ((b backtrack-pointer)
+      (wam-backtrack-pointer wam)))
+  (:returns trail-index)
+  (wam-stack-word wam (+ b 5)))
+
+(defun* wam-stack-choice-h
+    ((wam wam)
+     &optional
+     ((b backtrack-pointer)
+      (wam-backtrack-pointer wam)))
+  (:returns heap-index)
+  (wam-stack-word wam (+ b 6)))
+
 
-  "
-  ;; TODO: this sucks, fix it
-  (with-slots (stack) wam
-    (repeat words
-      (if (= +stack-limit+ (fill-pointer stack))
-        (error "WAM stack exhausted.")
-        (vector-push-extend 0 stack))))
-  (values))
+(defun* wam-stack-choice-arg
+    ((wam wam)
+     (n arity)
+     &optional
+     ((b backtrack-pointer)
+      (wam-backtrack-pointer wam)))
+  (:returns heap-index)
+  (wam-stack-word wam (+ b 7 n)))
+
+(defun (setf wam-stack-choice-arg)
+    (new-value wam n &optional (b (wam-backtrack-pointer wam)))
+  (setf (wam-stack-word wam (+ b 7 n))
+        new-value))
 
-(defun* wam-stack-pop-environment! ((wam wam))
-  "Pop an environment (stack frame) off the WAM stack."
-  (let ((frame-size (wam-stack-frame-size wam)))
-    (with-slots (stack environment-pointer) wam
-      (setf environment-pointer (wam-stack-frame-ce wam)) ; E <- CE
-      (decf (fill-pointer stack) frame-size)))) ; its fine
+(defun* wam-stack-choice-arg-cell
+    ((wam wam)
+     (n arity)
+     &optional
+     ((b backtrack-pointer)
+      (wam-backtrack-pointer wam)))
+  (:returns heap-cell)
+  (wam-heap-cell wam (wam-stack-choice-arg wam n b)))
+
+
+(defun* wam-stack-choice-size
+    ((wam wam)
+     &optional
+     ((b backtrack-pointer)
+      (wam-backtrack-pointer wam)))
+  (:returns stack-choice-size)
+  "Return the size of the choice frame starting at backtrack pointer `b`."
+  (+ (wam-stack-choice-n wam b) 7))
+
+
+(defun* wam-stack-pop-choice! ((wam wam))
+  "Pop a choice frame off the WAM stack."
+  (let ((size (wam-stack-choice-size wam)))
+    (with-slots (stack backtrack-pointer) wam
+      (setf backtrack-pointer
+            (wam-stack-choice-cb wam backtrack-pointer)) ; B <- CB
+      (decf (fill-pointer stack) size)))) ; its fine
 
 
 ;;;; Resetting
@@ -248,15 +411,23 @@
 (defun* wam-truncate-stack! ((wam wam))
   (setf (fill-pointer (wam-stack wam)) 0))
 
+(defun* wam-truncate-trail! ((wam wam))
+  (setf (fill-pointer (wam-trail wam)) 0))
+
+(defun* wam-truncate-unification-stack! ((wam wam))
+  (setf (fill-pointer (wam-unification-stack wam)) 0))
+
 (defun* wam-reset-local-registers! ((wam wam))
   (loop :for i :from 0 :below +register-count+ :do
         (setf (wam-local-register wam i)
               (1- +heap-limit+)))
-  (setf (wam-s wam) nil))
+  (setf (wam-subterm wam) nil))
 
 (defun* wam-reset! ((wam wam))
   (wam-truncate-heap! wam)
   (wam-truncate-stack! wam)
+  (wam-truncate-trail! wam)
+  (wam-truncate-unification-stack! wam)
   (wam-reset-local-registers! wam)
   (setf (wam-program-counter wam) 0
         (wam-continuation-pointer wam) 0
@@ -324,8 +495,11 @@
   (:returns (or null code-index))
   (gethash functor (wam-code-labels wam)))
 
-(defun (setf wam-code-label) (new-value wam functor)
-  (setf (gethash functor (wam-code-labels wam)) new-value))
+;; Note that this takes a functor/arity and not a cons.
+(defun (setf wam-code-label) (new-value wam functor arity)
+  (setf (gethash (wam-ensure-functor-index wam (cons functor arity))
+                 (wam-code-labels wam))
+        new-value))
 
 
 (defun* wam-load-query-code! ((wam wam) query-code)
@@ -380,7 +554,7 @@
   If S is unbound, throws an error.
 
   "
-  (let ((s (wam-s wam)))
+  (let ((s (wam-subterm wam)))
     (if (null s)
       (error "Cannot dereference unbound S register.")
       (wam-heap-cell wam s))))