join2

3 files changed, 62 insertions, 3 deletions

diff --git a/mcgoron.weight-balanced-trees.internal.scm b/mcgoron.weight-balanced-trees.internal.scm
index a3c9059..8d9dc6b 100644
--- a/mcgoron.weight-balanced-trees.internal.scm
+++ b/mcgoron.weight-balanced-trees.internal.scm
@@ -219,3 +219,20 @@
       ((negative? dir) (s-join data ("<" left) (">" right)))
       (else (wb-tree-node data left right)))))
 
+(: join2 (node-type node-type --> node-type))
+(define (join2 left right)
+  (define split-last
+    (the ((struct <wb-tree>) --> node-type *)
+         (lambda (tree)
+           (with-node (tree data ("<" left) (">" right))
+             (if (null? right)
+                 (values left data)
+                 (let-values (((new-right new-data)
+                              (split-last right)))
+                   (values (join data left new-right)
+                           new-data)))))))
+  (if (null? left)
+      right
+      (let-values (((new-left new-data)
+                    (split-last left)))
+        (join new-data new-left right))))
diff --git a/mcgoron.weight-balanced-trees.internal.sld b/mcgoron.weight-balanced-trees.internal.sld
index f5f5362..804a8bf 100644
--- a/mcgoron.weight-balanced-trees.internal.sld
+++ b/mcgoron.weight-balanced-trees.internal.sld
@@ -21,12 +21,15 @@
     (else (begin
             (define-syntax :
               (syntax-rules ()
-                ((: dummy ...) (begin)))))))
+                ((: dummy ...) (begin))))
+            (define-syntax the
+              (syntax-rules ()
+                ((the type expression) expression))))))
   (export wb-tree-node? non-null-wb-tree-node?
           get-data get-left get-right get-size
           balanced?
           in-order-vector->node
           node->in-order-list
-          join)
+          join join2)
   (include "mcgoron.weight-balanced-trees.internal.scm"))
 
diff --git a/tests/run.scm b/tests/run.scm
index a280233..1aea3b0 100644
--- a/tests/run.scm
+++ b/tests/run.scm
@@ -112,7 +112,12 @@
                                   (max-vector-length)))))
 
 ;;; ;;;;;;;;;
-;;; Test that join preserves the order of nodes
+;;; Make a "split" vector. A "split" vector is two sorted vectors and a 
+;;; number, where for all indicies i
+;;; 
+;;;      v1[i] < number
+;;;      number < v2[i]
+;;;   
 ;;; ;;;;;;;;;
 
 (define (split-in-order-vector vec)
@@ -131,6 +136,10 @@
          (make-in-order _)
          (gmap split-in-order-vector _)))
 
+;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;;; join
+;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
 (define (join-preserves-order-of vec)
   ;; Test if joining two in-order vectors with a node in between will
   ;; preserve order in the tree.
@@ -169,3 +178,33 @@
   (define (joined-vector-generator of)
     (gmap join-split-vectors (split-vector-generator of)))
   (test-property balanced? (list (joined-vector-generator (max-vector-length)))))
+
+;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;
+;;; join2
+;;; ;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+(define (join2-preserves-order-of vec)
+  (let* ((left (vector-ref vec 0))
+         (right (vector-ref vec 2))
+         (ret (node->in-order-list
+               (join2 (in-order-vector->node left)
+                      (in-order-vector->node right))))
+         (orig (append (vector->list left)
+                       (vector->list right))))
+    (equal? ret orig)))
+
+(test-group "join2 preserves order"
+  (test-property join2-preserves-order-of
+                 (list (split-vector-generator
+                        (max-vector-length)))))
+
+(test-group "join2 is balanced"
+  (define (join2-split-vector-generator)
+    (define (join2-of-split vec)
+      (join2 (in-order-vector->node (vector-ref vec 0))
+             (in-order-vector->node (vector-ref vec 2))))
+    (gmap join2-of-split (split-vector-generator
+                          (max-vector-length))))
+  (test-property balanced?
+                 (list (join2-split-vector-generator))))
+