module datatypes

pub struct Set[T] {
mut:
	elements map[T]u8
}

// checks the element is exists.
pub fn (set Set[T]) exists(element T) bool {
	return element in set.elements
}

// adds the element to set, if it is not present already.
pub fn (mut set Set[T]) add(element T) {
	set.elements[element] = 1
}

// removes the element from set.
pub fn (mut set Set[T]) remove(element T) {
	set.elements.delete(element)
}

// pick returns an arbitrary element of set, if set is not empty.
pub fn (set Set[T]) pick() !T {
	for k, _ in set.elements {
		return k
	}
	return error('Set is empty.')
}

// rest returns the set consisting of all elements except for the arbitrary element.
pub fn (mut set Set[T]) rest() ![]T {
	element := set.pick()!
	return set.elements.keys().filter(it != element)
}

// pop returns an arbitrary element and deleting it from set.
pub fn (mut set Set[T]) pop() !T {
	element := set.pick()!
	set.elements.delete(element)
	return element
}

// delete all elements of set.
pub fn (mut set Set[T]) clear() {
	set.elements = map[T]u8{}
}

// == checks whether the two given sets are equal (i.e. contain all and only the same elements).
pub fn (l Set[T]) == (r Set[T]) bool {
	if l.elements.len != r.elements.len {
		return false
	}
	for e, _ in r.elements {
		if e !in l.elements {
			return false
		}
	}
	return true
}

// is_empty checks whether the set is empty or not.
pub fn (set Set[T]) is_empty() bool {
	return set.size() == 0
}

// size returns the number of elements in the set.
pub fn (set Set[T]) size() int {
	return set.elements.len
}

// copy returns a copy of all the elements in the set.
pub fn (set Set[T]) copy() Set[T] {
	return Set[T]{
		elements: set.elements.clone()
	}
}

// add_all adds the whole `elements` array to the set
pub fn (mut set Set[T]) add_all(elements []T) {
	for element in elements {
		set.add(element)
	}
}

// @union returns the union of the two sets.
pub fn (l Set[T]) @union(r Set[T]) Set[T] {
	mut set := l.copy()
	for e, _ in r.elements {
		set.add(e)
	}
	return set
}

// intersection returns the intersection of sets.
pub fn (l Set[T]) intersection(r Set[T]) Set[T] {
	mut set := l.copy()
	for e, _ in l.elements {
		if !r.exists(e) {
			set.remove(e)
		}
	}
	for e, _ in r.elements {
		if !l.exists(e) {
			set.remove(e)
		}
	}
	return set
}

// - returns the difference of sets.
pub fn (l Set[T]) - (r Set[T]) Set[T] {
	mut set := l.copy()
	for e, _ in l.elements {
		if r.exists(e) {
			set.remove(e)
		}
	}
	return set
}

// subset returns true if the set `r` is a subset of the set `l`.
pub fn (l Set[T]) subset(r Set[T]) bool {
	for e, _ in r.elements {
		if e !in l.elements {
			return false
		}
	}
	return true
}

// array returns an array representation of the set
pub fn (l Set[T]) array() []T {
	return l.elements.keys()
}
