TreeSet / TreeMap implementation

So I was doing some work trying to manage tasks and got irritated that I didn’t have a decent map data type in AppleScript. I didn’t really want an extension, so I figured I’d just port something quickly.

I decided on a straightforward AA tree, it’s less rigorous than a red-black tree, but still quite performant. The code is pretty well lifted from the wikipedia page on AA trees :wink: A lot of people prefer hashing, but hash tables can perform very badly with some data sets, whereas trees will perform reasonably in all cases. And, of course, you get your data sorted for free, which is generally what people want.

You could probably take the script objects and save them using the normal means; I haven’t tried that.

Usage is pretty simple:


set theSet to makeSet()
theSet's add(5)
theSet's keys()
theSet's remove(5)
theSet's contains(5) is false
set theMap to makeMap()
theMap's add("key", "value")
theMap's keys() is {"key"}
theMap's values() is {"value"}
theMap's keyValuePairs() is {{"key", "value"}}
theMap's remove("key")
theMap's contains("key") is false
theMap's lookUp("key", "default") is "default"

No bulk methods like union or difference, I’m afraid.


on makeSet()
	script SortedSet
		property root : missing value
		on treeNode(the_key)
			return {key:the_key, level:1, leftN:missing value, rightN:missing value}
		end treeNode
		on copyData(S, D)
			set D's key to S's key
		end copyData
		on skew(T)
			if T is missing value then
				return missing value
			else if T's leftN is missing value then
				return T
			else if T's leftN's level is equal to T's level then
				-- Swap the pointers of horizontal leftN links.
				set L to T's leftN
				set T's leftN to L's rightN
				set L's rightN to T
				return L
			else
				return T
			end if
		end skew
		on split(T)
			if T is missing value then
				return missing value
			else if T's rightN is missing value or T's rightN's rightN is missing value then
				return T
			else if T's level is equal to T's rightN's rightN's level then
				local R
				set R to T's rightN
				set T's rightN to R's leftN
				set R's leftN to T
				set R's level to (R's level) + 1
				return R
			else
				return T
			end if
		end split
		on insert(X, T)
			-- Do the normal binary tree insertion procedure. Set the result of the
			-- recursive call to the correct child in case a new node was created or the
			-- root of the subtree changes.
			if T is missing value then
				-- Create a new leaf node with X.
				return X
			else if X's key < T's key then
				set T's leftN to insert(X, T's leftN)
			else if X's key > T's key then
				set T's rightN to insert(X, T's rightN)
			else if X's key = T's key then
				copyData(X, T)
			end if
			--  Perform skew and then split. The conditionals that determine whether or
			--  not a rotation will occur or not are inside of the procedures, as given
			--  above.
			set T to skew(T)
			set T to split(T)
			return T
		end insert
		on add(X)
			local R
			set R to insert(treeNode(X), root)
			set root to R
		end add
		on keys()
			script working
				property the_list : {}
				on handle(N)
					set the_list's end to N's key
				end handle
			end script
			if root is not missing value then
				traverse(root, working)
			end if
			return working's the_list
		end keys
		on traverse(N, obj)
			if N's leftN is not missing value then
				traverse(N's leftN, obj)
			end if
			obj's handle(N)
			if N's rightN is not missing value then
				traverse(N's rightN, obj)
			end if
		end traverse
		on min(A, B)
			if A < B then
				return A
			else
				return B
			end if
		end min
		on level_of(T)
			if T is missing value then
				return 0
			else
				return T's level
			end if
		end level_of
		on leaf(T)
			return T's leftN is missing value and T's rightN is missing value
		end leaf
		on decrease_level(T)
			local should_be
			set should_be to min(level_of(T's leftN), level_of(T's rightN)) + 1
			if should_be < T's level then
				set T's level to should_be
				if should_be < level_of(T's rightN) then
					set T's rightN's level to should_be
				end if
			end if
			return T
		end decrease_level
		on predecessor(T)
			local N, P
			set N to T's leftN
			if N's rightN is missing value then
				set T's leftN to missing value
				return N
			end if
			repeat while N's rightN is not missing value
				set P to N
				set N to N's rightN
			end repeat
			set P's rightN to missing value
			return N
		end predecessor
		on successor(T)
			local N, P
			set N to T's rightN
			if N's leftN is missing value then
				set T's rightN to missing value
				return N
			end if
			repeat while N's leftN is not missing value
				set P to N
				set N to N's leftN
			end repeat
			set P's leftN to missing value
			return N
		end successor
		on deleteNode(K, T)
			local L
			if T is missing value then
				return missing value
			else if K > T's key then
				set T's rightN to deleteNode(K, T's rightN)
			else if K < T's key then
				set T's leftN to deleteNode(K, T's leftN)
			else
				if leaf(T) then
					return missing value
				else if T's leftN is missing value then
					set L to successor(T)
					set T's rightN to deleteNode(L's key, T's rightN)
					copyData(L, T)
				else
					set L to predecessor(T)
					set T's leftN to deleteNode(L's key, T's leftN)
					copyData(L, T)
				end if
			end if
			--Rebalance the tree. Decrease the level of all nodes in this level if
			--necessary, and then skew and split all nodes in the new level.				
			set T to decrease_level(T)
			set T to skew(T)
			set T's rightN to skew(T's rightN)
			if T's rightN is not missing value then
				set T's rightN's rightN to skew(T's rightN's rightN)
			end if
			set T to split(T)
			set T's rightN to split(T's rightN)
			return T
		end deleteNode
		on remove(K)
			set root to deleteNode(K, root)
		end remove
		on nestedLists(T)
			local X
			set X to "["
			if T's leftN is not missing value then
				set X to X & nestedLists(T's leftN) & "|"
			end if
			set X to X & (T's key as string)
			if T's rightN is not missing value then
				set X to X & "|" & nestedLists(T's rightN)
			end if
			return X & "]"
		end nestedLists
		on nested()
			return nestedLists(root)
		end nested
		on findNode(K, T)
			if T is missing value then
				return missing value
			else if K < T's key then
				return findNode(K, T's leftN)
			else if K > T's key then
				return findNode(K, T's rightN)
			else
				return T
			end if
		end findNode
		on containsKey(K)
			return findNode(K, root) is not missing value
		end containsKey
	end script
	return SortedSet
end makeSet

on makeMap()
	local SortedSet
	set SortedSet to my makeSet()
	script SortedMap
		property parent : SortedSet
		on treeNode(the_key, the_val)
			return {key:the_key, value:the_val, level:1, leftN:missing value, rightN:missing value}
		end treeNode
		on copyData(S, D)
			set D's key to S's key
			set D's value to S's value
		end copyData
		on add(K, V)
			set my root to insert(treeNode(K, V), my root)
		end add
		on lookup(K, D)
			local N
			set N to findNode(K, root)
			if N is missing value then
				return D
			else
				return N's value
			end if
		end lookup
		on values()
			script working
				property the_list : {}
				on handle(N)
					set the_list's end to N's value
				end handle
			end script
			if my root is not missing value then
				my traverse(my root, working)
			end if
			return working's the_list
		end values
		on keyValuePairs()
			script working
				property the_list : {}
				on handle(N)
					set the_list's end to {N's key, N's value}
				end handle
			end script
			if my root is not missing value then
				my traverse(my root, working)
			end if
			return working's the_list
		end keyValuePairs
	end script
	return SortedMap
end makeMap