UNSLISP/set.scm

503 lines
19 KiB
Scheme

;;; Copyright (C) Peter McGoron 2024
;;; This program is free software: you can redistribute it and/or modify
;;; it under the terms of the GNU General Public License as published by
;;; the Free Software Foundation, version 3 of the License.
;;;
;;; This program is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
;;; GNU General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with this program. If not, see <https://www.gnu.org/licenses/>.
;;; Persistent AVL sets and maps using JOIN.
;;; ;;;;;;;;;;;;;;;;
;;; Nodes, direction
;;; ;;;;;;;;;;;;;;;;
;;; Returns the slot number for
;;; =: The node data
;;; h: the height
;;; <: the left child
;;; >: the right child
(define %set:accessor
(lambda (sym)
(cond
((eq? sym '=) 0)
((eq? sym 'h) 1)
((eq? sym '<) 2)
((eq? sym '>) 3)
(else (error "invalid direction")))))
;;; Gets data from node value given accessor symbol.
(define set:get
(lambda (t sym)
(vector-ref t (%set:accessor sym))))
(define %set->sexpr
(lambda (node)
(if (null? node)
'()
(list (list 'data (set:get node '=))
(list '< (%set->sexpr (set:get node '<)))
(list '> (%set->sexpr (set:get node '>)))))))
;;; Get the height of a node, handling the empty node.
(define %set:height
(lambda (node)
(if (null? node)
0
(set:get node 'h))))
;;; Get the difference between the heights of two trees.
(define %set:height-diff
(lambda (t1 t2)
(- (%set:height t1) (%set:height t2))))
;;; Get the balance factor of a tree.
(define %set:bal
(lambda (t) (%set:height-diff (set:get t '<)
(set:get t '>))))
;;; Set data in node given accessor symbol.
(define %set:set!
(lambda (node dir x)
(vector-set! node (%set:accessor dir) x)))
;;; Construct a new tree with data VAL.
(define %set:node
(lambda (val dir1 node1 dir2 node2)
(let ((node (vector val (+ 1
(max (%set:height node1)
(%set:height node2)))
'() '())))
(%set:set! node dir1 node1)
(%set:set! node dir2 node2)
node)))
(define set:node-new-val
(lambda (node newval)
(%set:node newval
'< (set:get node '<)
'> (set:get node '>))))
(define %set:invdir
(lambda (dir)
(cond
((eq? dir '<) '>)
((eq? dir '>) '<)
(else (error "invalid direction")))))
;;; ;;;;;;;;;;;;;;
;;; Tree rotations
;;; ;;;;;;;;;;;;;;
;;; Rotate NODE to the left (dir = '>) or right (dir = '<).
(define %set:rotate
(lambda (node dir)
(if (null? node)
#f
(let ((invdir (%set:invdir dir)))
(let ((child (set:get node invdir)))
(let ((to-swap (set:get child dir)))
(%set:node (set:get child '=)
dir (%set:node (set:get node '=)
dir (set:get node dir)
invdir to-swap)
invdir (set:get child invdir))))))))
;;; ;;;;;;;;;;;;;;;;;;;
;;; JOIN function for AVL trees.
;;; ;;;;;;;;;;;;;;;;;;;
;;; Handles rebalancing of the tree.
(define %set:join
(lambda (heavier val lighter heavier-dir)
(let ((heavy-val (set:get heavier '=))
(lighter-dir (%set:invdir heavier-dir)))
(let ((heavy-heavy (set:get heavier heavier-dir))
(heavy-light (set:get heavier lighter-dir)))
(if (<= (abs (%set:height-diff heavy-light lighter)) 1)
(let ((node (%set:node val
heavier-dir heavy-light
lighter-dir lighter)))
(if (<= (abs (%set:bal node)) 1)
(%set:node heavy-val
heavier-dir heavy-heavy
lighter-dir node)
(%set:rotate (%set:node heavy-val
heavier-dir heavy-heavy
lighter-dir
(%set:rotate node lighter-dir))
heavier-dir)))
(let ((new-light (%set:join heavy-light val lighter heavier-dir)))
(let ((node (%set:node heavy-val
heavier-dir heavy-heavy
lighter-dir new-light)))
(if (<= (abs (%set:bal node)) 1)
node
(%set:rotate node heavier-dir)))))))))
(define set:join
(lambda (val dir1 node1 dir2 node2)
(let ((diff (%set:height-diff node1 node2)))
(cond
((> diff 1) (%set:join node1 val node2 dir1))
((< diff -1) (%set:join node2 val node1 dir2))
(else (%set:node val dir1 node1 dir2 node2))))))
(define set:join2
(letrec
((join2
(lambda (left right)
(if (null? left)
right
(let ((split-last-tree (split-last left)))
(set:join (set:get split-last-tree '=)
'< (set:get split-last-tree '<)
'> right)))))
(split-last
(lambda (tree)
(let ((right (set:get tree '>)))
(if (null? right)
tree
(let ((last (split-last right)))
(%set:node (set:get last '=)
(set:join (set:get tree '=)
'< (set:get tree '<)
'> (set:get last '>)))))))))
join2))
;;; ;;;;;;;;;;;;;;;;;
;;; Element functions
;;; ;;;;;;;;;;;;;;;;;
;;; (SET:IN <=>) generates a search function for comparison function <=>.
;;; (SEARCH TREE DATA) searches TREE for a node that matches DATA.
;;; It will return the node that contains the matched DATA, or #F.
(define set:in
(lambda (<=>)
(lambda (tree data)
(letrec
((loop
(lambda (tree)
(if (null? tree)
#f
(let ((dir (<=> (set:get tree '=) data)))
(if (eq? dir '=)
tree
(loop (set:get tree dir))))))))
(loop tree)))))
;;; (SET:UPDATE <=>) generates an update function for <=>.
;;; (UPDATE TREE NODE UPDATE) inserts (UPDATE NODE #F) into the tree if
;;; no node comparing equal to NODE is found, and (UPDATE NODE OLD) if
;;; OLD compares equal to NODE.
(define set:update
(lambda (<=>)
(lambda (tree node update)
(letrec
((loop
(lambda (tree)
(if (null? tree)
(update node #f)
(let ((dir (<=> (set:get tree '=)
(set:get node '=))))
(if (eq? dir '=)
(update node tree)
(let ((invdir (%set:invdir dir)))
(set:join (set:get tree '=)
dir (loop (set:get tree dir))
invdir (set:get tree invdir)))))))))
(loop tree)))))
;;; (SET:INSERT <=>) generates an insert function for comparison function
;;; <=>.
;;; (INSERT TREE NODE) inserts NODE into TREE. It returns
;;; (CONS NEWTREE FOUND), where FOUND is the node that was replaced by
;;; NODE, and #F otherwise, and NEWTREE is the new root of the tree.
(define set:insert
(lambda (update)
(lambda (tree node)
(let ((found #f))
(let ((newroot (update tree node
(lambda (node oldnode)
(if (not oldnode)
node
(begin
(set! found oldnode)
(set:node-new-val oldnode
(set:get node '=))))))))
(cons newroot found))))))
;;; (SET:DELETE <=>) generates a delete function for comparison function
;;; <=>.
;;; (DELETE TREE DATA) deletes a node from TREE that compares equal to
;;; DATA. The function returns (CONS NEWTREE FOUND), where FOUND is the
;;; deleted node, or #F if not found, and NEWTREE is the root of the new
;;; tree.
(define set:delete
(lambda (<=>)
(lambda (tree node)
(let ((found #f))
(letrec
((loop
(lambda (tree)
(if (null? tree)
node
(let ((dir (<=> (set:get tree '=)
(set:get node '=))))
(if (eq? dir '=)
(begin
(set! found tree)
(set:join2 (set:get tree '<)
(set:get tree '>)))
(let ((invdir (%set:invdir dir)))
(set:join (set:get tree '=)
dir (loop (set:get tree dir))
invdir (set:get tree invdir)))))))))
(let ((newtree (loop tree)))
(cons newtree found)))))))
;;; ;;;;;;;;;;;;;;;;;;;;;;;
;;; Converting sets to maps
;;; ;;;;;;;;;;;;;;;;;;;;;;;
;;; Convert a <=> for sets to one for maps.
(define set:<=>-to-map
(lambda (<=>)
(lambda (x y)
(<=> (car x) (car y)))))
(define map:node-new-val
(lambda (node newkey newval)
(set:node-new-val node
(cons newkey newval))))
(define map:empty-node
(lambda (key val)
(%set:node (cons key val)
'< '() '> '())))
(define map:key
(lambda (node)
(car (set:get node '=))))
(define map:val
(lambda (node)
(cdr (set:get node '=))))
;;; ;;;;;;;;;;;
;;; For strings
;;; ;;;;;;;;;;;
(cond-expand
((not miniscm-unslisp)
(define (string<=> x y)
(cond
((string<? x y) '<)
((string>? x y) '>)
(else '=)))))
(cond-expand
((and (not miniscm-unslisp) (not r7rs))
(define (list-set! lst n val)
(if (= n 0)
(set-car! lst val)
(list-set! (cdr lst) (- n 1) val)))))
(define map:string<=> (set:<=>-to-map string<=>))
(define %smap:update (set:update map:string<=>))
(define %smap:insert (set:insert %smap:update))
;;; (SMAP:UPDATE TREE KEY UPDATE) updates NODE in TREE with
;;; (UPDATE NODE-WITH-KEY NODE), where NODE-WITH-KEY is an empty node
;;; with the key KEY, and (UPDATE NODE-WITH-KEY #F) if no NODE is
;;; found.
(define smap:update
(lambda (tree key update)
(%smap:update tree (map:empty-node key '()) update)))
;;; (SMAP:INSERT TREE KEY VAL) inserts (CONS KEY VAL) into TREE, and
;;; returns (CONS NEWROOT FOUND), where NEWROOT is the new root of
;;; the tree, and FOUND is #F if no element matching KEY was found,
;;; or the matching element if found.
(define smap:insert
(lambda (tree key val)
(%smap:insert tree (map:empty-node key val))))
(define %smap:search (set:in map:string<=>))
;;; (SMAP:SEARCH TREE KEY)
(define smap:search
(lambda (tree key)
(%smap:search tree (cons key '()))))
(define %smap:delete (set:delete map:string<=>))
(define smap:delete
(lambda (tree key)
(%smap:delete tree (map:empty-node key '()))))
;;; ;;;;;
;;; Tests
;;; ;;;;;
;;; LST is a list of elements of the form
;;; (KEY VAL ALREADY-IN)
;;; where ALREADY-IN is #F for an element not in the set, or the value
;;; that should be in the set.
(define %set:operate-all
(lambda (f tree lst)
(if (null? lst)
tree
(let ((key (list-ref (car lst) 0))
(val (list-ref (car lst) 1))
(already-in (list-ref (car lst) 2)))
(let ((insert-return (f tree key val)))
(cond
((and already-in (not (cdr insert-return)))
"should have been found")
((and already-in (not (equal? already-in
(map:val (cdr insert-return)))))
"found is not correct")
(else (%set:operate-all f (car insert-return) (cdr lst)))))))))
(define %set:insert-all
(lambda (tree lst)
(%set:operate-all smap:insert tree lst)))
(define %set:search-all
(lambda (tree lst)
(%set:operate-all (lambda (tree key _)
(let ((search-res (smap:search tree key)))
(cons tree search-res))) tree lst)))
(define %set:delete-all
(lambda (tree lst)
(%set:operate-all (lambda (tree key _)
(smap:delete tree key)) tree lst)))
(define %set:tests
(list
(cons "rotate right"
(lambda ()
(let ((right (%set:rotate (%set:node 1
'< (%set:node 2
'< (%set:node 3
'< '()
'> '())
'> (%set:node 4
'< '()
'> '()))
'> (%set:node 5 '< '() '> '()))
'>)))
(cond
((not (eqv? (set:get right '=) 2)) "bad parent")
((not (eqv? (set:get (set:get right '>) '=) 1)) "bad right child")
((not (eqv? (set:get (set:get right '<) '=) 3)) "bad left child")
((not (eqv? (set:get (set:get (set:get right '>) '>) '=) 5))
"bad right child of right child")
((not (eqv? (set:get (set:get (set:get right '>) '<) '=) 4))
"bad left child of right child")
(else #t)))))
(cons "rotate left"
(lambda ()
(let ((right (%set:rotate (%set:node 1
'> (%set:node 2
'< (%set:node 3
'< '()
'> '())
'> (%set:node 4
'< '()
'> '()))
'< (%set:node 5 '< '() '> '()))
'<)))
(cond
((not (eqv? (set:get right '=) 2)) "bad parent")
((not (eqv? (set:get (set:get right '>) '=) 4)) "bad right child")
((not (eqv? (set:get (set:get right '<) '=) 1)) "bad left child")
((not (eqv? (set:get (set:get (set:get right '<) '>) '=) 3))
"bad right child of left child")
((not (eqv? (set:get (set:get (set:get right '<) '<) '=) 5))
"bad left child of left child")
(else #t)))))
(cons "insert then delete"
(lambda ()
(let ((insert-return (smap:insert '() (string #\a) 5)))
(cond
((not (pair? insert-return)) "invalid insert return")
((cdr insert-return) "string found in empty tree")
(else
(let ((tree (car insert-return)))
(let ((found (smap:search tree (string #\a))))
(cond
((not found) "string not in tree")
((not (equal? (map:key tree) (string #\a)))
"returned key not equal to a")
((not (equal? (map:val tree) 5))
"returned value not equal to 5")
(else
(let ((delete-return (smap:delete tree (string #\a))))
(cond
((not (pair? delete-return))
"invalid delete return")
((not (cdr delete-return)) "string not found")
((not (eqv? (car delete-return) '()))
"returned tree not null")
(else #t))))))))))))
(cons "insert a few unique then delete"
(lambda ()
(let ((to-insert (list
(list (string #\a #\b #\c) 1 #f)
(list (string #\a #\b #\d) 2 #f)
(list (string #\d #\e #\f) 3 #f)
(list (string #\1 #\2 #\3 #\a #\C) 4 #f))))
(display "insert all") (newline)
(let ((tree (%set:insert-all '() to-insert)))
(if (string? tree)
tree
(begin
(for-each (lambda (x)
(list-set! x 2 (list-ref x 1)))
to-insert)
(display "search all") (newline)
(let ((res (%set:search-all tree to-insert)))
(if (string? res)
res
(begin
(display "delete all") (newline)
(let ((tree (%set:delete-all tree to-insert)))
(cond
((string? tree) tree)
((not (null? tree)) "did not delete everything")
(else #t))))))))))))
(cons "insert a few, update"
(lambda ()
(let ((tree (%set:insert-all '()
(list
(list (string #\a #\b #\c #\d) 1 #f)
(list (string #\e #\f #\g #\h) 2 #f)
(list (string #\1 #\4 #\2 #\9 #\3) 3 #f)
(list (string #\a #\b #\c #\d #\e) 4 #f)))))
(if (string? tree)
tree
(let ((tree (smap:update tree
(string #\a #\b #\c #\d #\e)
(lambda (node oldnode)
(map:node-new-val oldnode
(string #\a #\b #\c #\d #\e)
10)))))
(let ((res (%set:search-all tree
(list
(list (string #\a #\b #\c #\d) 1 1)
(list (string #\e #\f #\g #\h) 2 2)
(list (string #\1 #\4 #\2 #\9 #\3) 3 3)
(list (string #\a #\b #\c #\d #\e) 10 10)))))
(if (string? res)
res
#t)))))))))