internal/pkg/layout/linear_map.go

// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package layout

import (
	"fmt"
	"sort"
)

// LinearMap manages a sequence of contiguous elements stacked along one dimension.
// It provides efficient O(log N) lookup for coordinates and intersections.
//
// Internal Representation:
// It uses an N+1 boundary array where boundaries[i] is the starting offset
// of element i, and boundaries[N] is the total extent. This guarantees
// contiguity by design.
type LinearMap struct {
	// boundaries stores the starting coordinates of each element, plus the end
	// coordinate of the last element. len(boundaries) == N + 1.
	//
	// Invariant: The array is monotonically non-decreasing (boundaries[i] <= boundaries[i+1]).
	// This is guaranteed because element sizes are required to be non-negative.
	boundaries []int
}

// NewLinearMap constructs a map from a slice of sizes, calculating
// the N+1 boundaries automatically starting from the local origin (0).
//
// It panics if any size is negative, as this represents a logic error in
// the calling layout code.
func NewLinearMap(sizes []int) *LinearMap {
	count := len(sizes)
	b := make([]int, count+1)
	current := 0
	for i, s := range sizes {
		if s < 0 {
			panic(fmt.Sprintf("layout: element %d has negative size %d", i, s))
		}
		b[i] = current
		current += s
	}
	b[count] = current
	return &LinearMap{
		boundaries: b,
	}
}

// Count returns the number of elements in the map.
func (m *LinearMap) Count() int {
	if len(m.boundaries) == 0 {
		return 0
	}
	return len(m.boundaries) - 1
}

// Total returns the full horizontal or vertical extent of the map.
func (m *LinearMap) Total() int {
	if len(m.boundaries) == 0 {
		return 0
	}
	return m.boundaries[len(m.boundaries)-1]
}

// Span returns the spatial footprint of the i-th element.
func (m *LinearMap) Span(i int) Span {
	if i < 0 || i >= m.Count() {
		return Span{}
	}
	return Span{
		Offset: m.boundaries[i],
		Size:   m.boundaries[i+1] - m.boundaries[i],
	}
}

// At returns the index of the element containing the given coordinate.
//
// Example:
//
//	Sizes: [10, 20, 30] -> Boundaries: [0, 10, 30, 60]
//	At(5)  -> 0 (Inside [0, 10))
//	At(15) -> 1 (Inside [10, 30))
//	At(40) -> 2 (Inside [30, 60))
//
// Returns -1 if the coordinate is out of bounds.
func (m *LinearMap) At(coord int) int {
	if coord < 0 || coord >= m.Total() {
		return -1
	}

	// Search for the largest boundary <= coord.
	// We look for the first boundary that is strictly greater than coord.
	idx := sort.Search(len(m.boundaries), func(i int) bool {
		return m.boundaries[i] > coord
	})
	return idx - 1
}

// Intersects returns the range of indices that overlap with the given span.
//
// Example:
//
//	Sizes: [10, 10, 10] -> Boundaries: [0, 10, 20, 30]
//	Intersects({5, 10})  -> {0, 2} (Window [5, 15) overlaps items [0, 10) and [10, 20))
//	Intersects({25, 5})  -> {2, 3} (Window [25, 30) overlaps item [20, 30))
//	Intersects({35, 10}) -> {3, 3} (Window [35, 45) past end, empty range)
func (m *LinearMap) Intersects(window Span) Interval {
	if window.IsEmpty() || m.Count() == 0 {
		return Interval{0, 0}
	}

	// 1. Find the first element that ends after the window starts.
	startIdx := sort.Search(m.Count(), func(i int) bool {
		// Element i ends at boundaries[i+1].
		return m.boundaries[i+1] > window.Offset
	})

	if startIdx == m.Count() {
		return Interval{startIdx, startIdx}
	}

	// 2. Find the first element that starts at or after the window ends.
	// All elements before this one intersect the window.
	if window.IsUnbounded() {
		return Interval{startIdx, m.Count()}
	}

	endIdx := sort.Search(m.Count(), func(i int) bool {
		return m.boundaries[i] >= window.End()
	})

	return Interval{startIdx, endIdx}
}