<commit_before>package trie

import (
	"sync"

	"github.com/jellevandenhooff/keytree/crypto"
)

// A Dedup helps store only one of multiple identical trie nodes.  A Dedup is
// threadsafe.
type Dedup struct {
	mu    sync.Mutex
	nodes map[crypto.Hash]dedupInfo
}

func NewDedup() *Dedup {
	return &Dedup{
		nodes: make(map[crypto.Hash]dedupInfo),
	}
}

type dedupInfo struct {
	refs int
	node *Node
}

func (d *Dedup) add(node *Node) *Node {
	if node == nil {
		return node
	}

	entry, found := d.nodes[node.Hash()]
	if found {
		entry.refs += 1
		return entry.node
	}

	node.Children[0] = d.add(node.Children[0])
	node.Children[1] = d.add(node.Children[1])
	d.nodes[node.Hash()] = dedupInfo{
		refs: 1,
		node: node,
	}
	return node
}

func (d *Dedup) Add(node *Node) *Node {
	d.mu.Lock()
	defer d.mu.Unlock()

	return d.add(node)
}

func (d *Dedup) AddWithChildrenAlreadyAdded(node *Node) *Node {
	if node == nil {
		return node
	}

	d.mu.Lock()
	defer d.mu.Unlock()

	entry, found := d.nodes[node.Hash()]
	if found {
		d.remove(node.Children[0])
		d.remove(node.Children[1])
		entry.refs += 1
		return entry.node
	}

	d.nodes[node.Hash()] = dedupInfo{
		refs: 1,
		node: node,
	}
	return node
}

func (d *Dedup) findAndAdd(hash crypto.Hash, count int) *Node {
	d.mu.Lock()
	defer d.mu.Unlock()

	info, found := d.nodes[hash]
	if found {
		info.refs += count
		return info.node
	} else {
		return nil
	}
}

func (d *Dedup) FindAndAdd(hash crypto.Hash) *Node {
	return d.findAndAdd(hash, 1)
}

func (d *Dedup) FindAndDoNotAdd(hash crypto.Hash) *Node {
	return d.findAndAdd(hash, 0)
}

func (d *Dedup) remove(node *Node) {
	if node == nil {
		return
	}

	entry := d.nodes[node.Hash()]
	entry.refs -= 1

	if entry.refs == 0 {
		delete(d.nodes, node.Hash())
		d.remove(node.Children[0])
		d.remove(node.Children[1])
	}
}<commit_msg>hack dedup to cache everything...<commit_after>package trie

import (
	"sync"

	"github.com/jellevandenhooff/keytree/crypto"
)

// A Dedup helps store only one of multiple identical trie nodes.  A Dedup is
// threadsafe.
type Dedup struct {
	mu    sync.Mutex
	nodes map[crypto.Hash]dedupInfo
}

func NewDedup() *Dedup {
	return &Dedup{
		nodes: make(map[crypto.Hash]dedupInfo),
	}
}

type dedupInfo struct {
	refs int
	node *Node
}

func (d *Dedup) add(node *Node) *Node {
	if node == nil {
		return node
	}

	entry, found := d.nodes[node.Hash()]
	if found {
		entry.refs += 1
		return entry.node
	}

	node.Children[0] = d.add(node.Children[0])
	node.Children[1] = d.add(node.Children[1])
	d.nodes[node.Hash()] = dedupInfo{
		refs: 1,
		node: node,
	}
	return node
}

func (d *Dedup) Add(node *Node) *Node {
	d.mu.Lock()
	defer d.mu.Unlock()

	return d.add(node)
}

func (d *Dedup) AddWithChildrenAlreadyAdded(node *Node) *Node {
	if node == nil {
		return node
	}

	d.mu.Lock()
	defer d.mu.Unlock()

	entry, found := d.nodes[node.Hash()]
	if found {
		d.remove(node.Children[0])
		d.remove(node.Children[1])
		entry.refs += 1
		return entry.node
	}

	d.nodes[node.Hash()] = dedupInfo{
		refs: 1,
		node: node,
	}
	return node
}

func (d *Dedup) findAndAdd(hash crypto.Hash, count int) *Node {
	d.mu.Lock()
	defer d.mu.Unlock()

	info, found := d.nodes[hash]
	if found {
		info.refs += count
		return info.node
	} else {
		return nil
	}
}

func (d *Dedup) FindAndAdd(hash crypto.Hash) *Node {
	return d.findAndAdd(hash, 1)
}

func (d *Dedup) FindAndDoNotAdd(hash crypto.Hash) *Node {
	return d.findAndAdd(hash, 0)
}

func (d *Dedup) remove(node *Node) {
	return

	if node == nil {
		return
	}

	entry := d.nodes[node.Hash()]
	entry.refs -= 1

	if entry.refs == 0 {
		delete(d.nodes, node.Hash())
		d.remove(node.Children[0])
		d.remove(node.Children[1])
	}
}
{'fraction_non_alphanumeric': 0.15812313534038513, 'fraction_numerical': 0.008679142934635205, 'mean_word_length': 3.4062126642771804, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6908080', 'n_tokens_mistral': 1620, 'n_tokens_neox': 1533, 'n_words': 361}