package lattice

func (t *Tree) NeedsCompact() bool {
	return len(t.Pending) > 0 || t.FreeCount > t.nCount/8
}

func (t *Tree) Compact() int {
	for i := range t.Pending {
		t.propagateOne(t.Pending[i])
	}
	t.Pending = t.Pending[:0]

	if t.FreeCount == 0 {
		return 0
	}

	moves := map[uint32]uint32{}
	tail := t.nCount - 1

	freeSet := map[uint32]bool{}
	cur := t.FreeHead
	for cur != NullIndex {
		freeSet[cur] = true
		next := t.getNode(cur).Children[0]
		cur = next
	}

	for tail > 0 && freeSet[tail] {
		delete(freeSet, tail)
		tail--
	}

	for gap := range freeSet {
		if gap >= tail {
			continue
		}
		for tail > gap && freeSet[tail] {
			delete(freeSet, tail)
			tail--
		}
		if tail <= gap {
			break
		}
		*t.getNode(gap) = *t.getNode(tail)
		t.markNodeDirty(gap)
		moves[tail] = gap
		freeSet[tail] = true
		tail--
		for tail > 0 && freeSet[tail] {
			delete(freeSet, tail)
			tail--
		}
	}

	newCount := tail + 1
	for i := uint32(0); i < newCount; i++ {
		node := t.getNode(i)
		dirty := false
		for ci := 0; ci < 4; ci++ {
			if newIdx, ok := moves[node.Children[ci]]; ok {
				node.Children[ci] = newIdx
				dirty = true
			}
		}
		if dirty {
			t.markNodeDirty(i)
		}
	}
	for ri := 0; ri < 3; ri++ {
		if newIdx, ok := moves[t.Roots[ri]]; ok {
			t.Roots[ri] = newIdx
		}
	}

	reclaimed := int(t.FreeCount)
	if t.cache == nil {
		t.Nodes = t.Nodes[:newCount]
	}
	t.nCount = newCount
	t.FreeHead = NullIndex
	t.FreeCount = 0
	return reclaimed
}

func (t *Tree) propagateOne(pe PendingExp) {
	root := t.Roots[pe.Branch]
	t.propagateSubtree(root, pe.Factor)
}

func (t *Tree) propagateSubtree(idx uint32, factor uint8) {
	if idx == NullIndex {
		return
	}
	node := t.getNode(idx)
	if node.IsDeleted() {
		return
	}
	if node.Mult > 1 {
		node.Mult = 1
		t.markNodeDirty(idx)
	}
	for ci := 0; ci <= node.KeyCount(); ci++ {
		if node.Children[ci] != NullIndex {
			t.propagateSubtree(node.Children[ci], factor)
		}
	}
}