/* Cifre Standard Library
   (c) Jeremy Tregunna, 2006, All Rights Reserved.

   This program is open source software, and redistributable under the terms
   and conditions of the BSD Licence. */

Dictionary := SortedCollection clone do(
	Node := Object clone do(
		newSlot("key")
		newSlot("value")
		newSlot("left")
		newSlot("right")
		newSlot("parent")
		newSlot("isRed", true)

		with := method(aKey, aValue,
			if(call message arguments size != 2,
				InvalidArgumentException raise("'with' requires a key and a value as arguments only")
			)
			self clone setKey(aKey) setValue(aValue)
		)
	)
	Node removeSlot("type") // setSlotWithType related -- we don't want a type slot

	newSlot("root")
	newSlot("sentinel")
	newSlot("lastNodeFound")
	newSlot("size", 0)

	init := method(
		node := Node clone
		setSentinel(node)
		sentinel setLeft(sentinel)
		sentinel setRight(sentinel)
		setLastNodeFound(sentinel)
		setRoot(sentinel)
	)

	docSlot("atPut(aKey, aValue)", "Stores the object aValue in the Dictionary associated with the key aKey. Returns the modified Dictionary.")
	atPut := method(aKey, aValue,
		if(call message arguments size < 1,
			InvalidArgumentException raise("atPut requires at least 1 argument")
		)
		checkKeyValidity(aKey)
		result := 0
		node := Node clone
		temp := root
		while(temp != sentinel,
			node setParent(temp)
			result = cmp(aKey, temp key)
			if(result == 0, return node)
			if(result > 0, temp = temp right, temp = temp left)
		)
		node setKey(aKey)
		node setValue(aValue)
		node setLeft(self sentinel)
		node setRight(self sentinel)
		if(node parent,
			result = cmp(node key, node parent key)
			if(result > 0, node parent setRight(node), node parent setLeft(node))
		,
			setRoot(node)
		)
		restoreAfterInsert(node)
		self setSize(self size + 1)
		self
	)

	docSlot("at(aKey)", "Retrieves the object associated with the key aKey, and returns it. Raises a NotFoundException if the object was not found.")
	at := method(aKey,
		checkKeyValidity(aKey)
		result := 0
		if(cmp(aKey, lastNodeFound ?key) == 0,
			return lastNodeFound value
		,
			node := root
			while(node != sentinel,
				result = cmp(aKey, node key)
				if(result == 0,
					setLastNodeFound(node)
					return node value
				)
				if(result < 0, node = node left, node = node right)
			)
		)
		NotFoundException raise("item not found")
	)

	docSlot("atIfAbsentDo(aKey, msg)", "Retrieves the object associated with the key aKey, and returns it. If the key is not found, executes the message tree msg, and returns its result.")
	atIfAbsentDo := method(aKey,
		r := nil
		try(r = at(aKey)) catch(NotFoundException, r = call evalArgAt(1))
		r
	)

	docSlot("atIfAbsentPut(aKey, msg)", "Searches the Dictionary for the key aKey. If the key is not found, evaluates the message msg within the context of the caller, and puts the value returned into the Dictionary associated with the key aKey. Returns the item at the key, or the result of the message.")
	atIfAbsentPut := method(aKey,
		r := nil
		try(r = at(aKey)) catch(NotFoundException, r = atPut(aKey, call evalArgAt(1)))
		r
	)

	docSlot("removeAt(aKey)", "Removes the key->value pair from the dictionary associated with aKey. If the key is not found, a NotFoundException is raised.")
	removeAt := method(aKey,
		checkKeyValidity(aKey)
		node := nil
		if(cmp(aKey, lastNodeFound ?key) == 0,
			node = lastNodeFound
		,
			node = root
			while(node != sentinel,
				result := cmp(aKey, node key)
				if(result == 0, break)
				if(result < 0, node = node left, node = node right)
			)
			if(node == sentinel, NotFoundException raise("item not found"))
		)
		delete(node)
		self setSize(self size - 1)
		node
	)

	docSlot("empty", "Clears all items in the dictionary. Returns the emptied dictionary.")
	empty := method(
		setRoot(sentinel)
		setSize(0)
		self
	)

	docSlot("isEmpty", "Returns true if the dictionary is empty, false otherwise.")
	isEmpty := method(root isNil or root == sentinel)

	docSlot("hasKey(aKey)", "Returns true if the key aKey is present in the receiver, or false otherwise.")
	hasKey := method(aKey,
		checkKeyValidity(aKey)
		r := true
		try(at(aKey)) catch(NotFoundException, r = false)
		r
	)

	docSlot("hasValue(aValue)", "Returns true if the object aValue is present in the receiver, or false otherwise.")
	hasValue := method(aValue,
		r := false
		if(self isEmpty not,
			self foreach(v, if(getSlot("v") == getSlot("aValue"), r = true))
		)
		r
	)

	docSlot("keys", "Returns a LinkedList of the receivers keys.")
	keys := method(
		r := LinkedList clone
		if(size > 0, self foreach(k, v, r append(getSlot("k"))))
		r
	)

	docSlot("values", "Returns a LinkedList of the receivers values.")
	values := method(
		r := LinkedList clone
		if(size > 0, self foreach(v, r append(getSlot("v"))))
		r
	)

	docSlot("foreach([aKey], aValue, msg)", "For each key/value pair, sets the locals aKey to the key, and aValue to the value and executes the message msg.")
	foreach := method(
		r := nil
		key := call argAt(0)
		val := call argAt(1)
		body := call argAt(2)
		if(val isNil, InvalidArgumentException raise("'foreach' requires at least 2 arguments."))
		if(body isNil,
			key = nil
			val = call argAt(0)
			body = call argAt(1)
		)
		traverse(root, Stack clone) foreach(node,
			if(getSlot("key"), call sender setSlot(key name, node key))
			call sender setSlot(val name, node value)
			call sender doMessage(body)
		)
	)

	docSlot("asString", "Returns a string containing the contents of the dictionary.")
	asString := method(
		s := Sequence with(self type asMutable lowercase .. "(")
		if(self size != 0,
			i := 0
			self foreach(k, v,
				s appendSeq(k asSimpleString .. " : " .. getSlot("v") asSimpleString)
				if(i != self size - 1, s appendSeq(", "))
				i = i + 1
			)
		)
		s appendSeq(")") asSymbol
	)

	// Undocumented routines

	restoreAfterInsert := method(x,
		y := nil
		while((x != root) and x parent isRed,
			if(x parent == x parent parent left,
				y = x parent parent right
				if(y and y isRed,
					x parent setIsRed(false)
					y setIsRed(false)
					x parent parent setIsRed(true)
					x = x parent parent
				,
					if(x == x parent right,
						x = x parent
						rotateLeft(x)
					)
					x parent setIsRed(false)
					x parent parent setIsRed(true)
					rotateRight(x parent parent)
				)
			,
				y = x parent parent left
				if(y and y isRed,
					x parent setIsRed(false)
					y setIsRed(false)
					x parent parent setIsRed(true)
					x = x parent parent
				,
					if(x == x parent left,
						x = x parent
						rotateRight(x)
					)
					x parent setIsRed(false)
					x parent parent setIsRed(true)
					rotateLeft(x parent parent)
				)
			)
		)
		root setIsRed(false)
		self
	)

	rotateLeft := method(x,
		y := x right
		x setRight(y left)
		if(y left != sentinel, y left setParent(x))
		if(y != sentinel, y setParent(x parent))
		if(x parent,
			if(x == x parent left, x parent setLeft(y), x parent setRight(y))
		,
			setRoot(y)
		)
		y setLeft(x)
		if(x != sentinel, x setParent(y))
		self
	)

	rotateRight := method(x,
		y := x left
		x setLeft(y right)
		if(y right != sentinel, y right setParent(x))
		if(y != sentinel, y setParent(x parent))
		if(x parent,
			if(x == x parent right, x parent setRight(y), x parent setLeft(y))
		,
			setRoot(y)
		)
		y setRight(x)
		if(x != sentinel, x setParent(y))
		self
	)

	delete := method(z,
		x := Node clone
		y := nil
		if((z left == sentinel) or z right == sentinel,
			y = z
		,
			y = z right
			while(y left != sentinel, y = y left)
		)
		if(y left != sentinel, x = y left, x = y right)
		x setParent(y parent)
		if(y parent,
			if(y == y parent left, y parent setLeft(x), y parent setRight(x))
		,
			setRoot(x)
		)
		if(y != z,
			z setKey(y key)
			z setValue(y value)
		)
		if(y isRed not, restoreAfterDelete(x))
		setLastNodeFound(sentinel)
		self
	)

	restoreAfterDelete := method(x,
		y := nil
		while((x != root) and x isRed not,
			if(x == x parent left,
				y = x parent right
				if(y isRed,
					y setIsRed(false)
					x parent setIsRed(true)
					rotateLeft(x parent)
					y = x parent right
				)
				if(y left isRed not and y right isRed not,
					y setIsRed(true)
					x = x parent
				,
					if(y right isRed not,
						y left setIsRed(false)
						y setIsRed(true)
						rotateRight(y)
						y = x parent right
					)
					y setIsRed(x parent isRed)
					x parent setIsRed(false)
					y right setIsRed(false)
					rotateLeft(x parent)
					x = root
				)
			,
				y = x parent left
				if(y isRed,
					y setIsRed(false)
					x parent setIsRed(true)
					rotateRight(x parent)
					y = x parent left
				)
				if(y right isRed not and y left isRed not,
					y setIsRed(true)
					x = x parent
				,
					if(y left isRed not,
						y right setIsRed(false)
						y setIsRed(true)
						rotateLeft(y)
						y = x parent left
					)
					y setIsRed(x parent isRed)
					x parent setIsRed(false)
					y left setIsRed(false)
					rotateRight(x parent)
					x = root
				)
			)
		)
		x setIsRed(false)
		self
	)

	traverse := method(node, stack,
		if(node != node left,
			if(node right and node right != sentinel,
				traverse(node left, stack)
			)
			if(stack and node key, stack push(node))
			if(node right and node right != sentinel,
				traverse(node right, stack)
			)
		)
		stack
	)

	checkKeyValidity := method(aKey,
		unless(aKey isKindOf(Sequence),
			InvalidArgumentException raise("key must be a Sequence")
		)
		if(aKey isMutable, aKey = aKey asSymbol)
	)
)