Test on real data

Move same generic parameters to associated types.
Work around private fields of SumResult<T>

Cargo.lock

@@ -635,15 +635,6 @@ version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
-[[package]]
-name = "bitmaps"
-version = "2.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "031043d04099746d8db04daf1fa424b2bc8bd69d92b25962dcde24da39ab64a2"
-dependencies = [
- "typenum",
-]
-
 [[package]]
 name = "block-buffer"
 version = "0.10.3"
@@ -1835,20 +1826,6 @@ dependencies = [
  "unicode-normalization",
 ]
 
-[[package]]
-name = "im"
-version = "15.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d0acd33ff0285af998aaf9b57342af478078f53492322fafc47450e09397e0e9"
-dependencies = [
- "bitmaps",
- "rand_core",
- "rand_xoshiro",
- "sized-chunks",
- "typenum",
- "version_check",
-]
-
 [[package]]
 name = "indexmap"
 version = "1.9.1"
@@ -2360,13 +2337,13 @@ dependencies = [
  "humantime",
  "humantime-serde",
  "hyper",
- "im",
  "itertools",
  "metrics",
  "nix 0.25.0",
  "num-traits",
  "once_cell",
  "pageserver_api",
+ "persistent_range_query",
  "pin-project-lite",
  "postgres",
  "postgres-protocol",
@@ -2485,6 +2462,14 @@ version = "2.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "478c572c3d73181ff3c2539045f6eb99e5491218eae919370993b890cdbdd98e"
 
+[[package]]
+name = "persistent_range_query"
+version = "0.1.0"
+dependencies = [
+ "rand",
+ "workspace_hack",
+]
+
 [[package]]
 name = "petgraph"
 version = "0.6.2"
@@ -2987,15 +2972,6 @@ dependencies = [
  "rand_core",
 ]
 
-[[package]]
-name = "rand_xoshiro"
-version = "0.6.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6f97cdb2a36ed4183de61b2f824cc45c9f1037f28afe0a322e9fff4c108b5aaa"
-dependencies = [
- "rand_core",
-]
-
 [[package]]
 name = "rayon"
 version = "1.5.3"
@@ -3630,16 +3606,6 @@ version = "0.3.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7bd3e3206899af3f8b12af284fafc038cc1dc2b41d1b89dd17297221c5d225de"
 
-[[package]]
-name = "sized-chunks"
-version = "0.6.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "16d69225bde7a69b235da73377861095455d298f2b970996eec25ddbb42b3d1e"
-dependencies = [
- "bitmaps",
- "typenum",
-]
-
 [[package]]
 name = "slab"
 version = "0.4.7"

libs/persistent_range_query/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "persistent_range_query"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+
+[dev-dependencies]
+rand = "0.8.3"

libs/persistent_range_query/src/lib.rs

@@ -0,0 +1,78 @@
+use std::ops::Range;
+
+pub mod naive;
+pub mod ops;
+pub mod segment_tree;
+
+/// Should be a monoid:
+/// * Identity element: for all a: combine(new_for_empty_range(), a) = combine(a, new_for_empty_range()) = a
+/// * Associativity: for all a, b, c: combine(combine(a, b), c) == combine(a, combine(b, c))
+pub trait RangeQueryResult<Key>: Sized + Clone {
+    // Clone is equivalent to combine with an empty range.
+
+    fn new_for_empty_range() -> Self;
+
+    // Contract: left_range.end == right_range.start
+    // left_range.start == left_range.end == right_range.start == right_range.end is still possible
+    fn combine(
+        left: &Self,
+        left_range: &Range<Key>,
+        right: &Self,
+        right_range: &Range<Key>,
+    ) -> Self;
+
+    fn add(left: &mut Self, left_range: &Range<Key>, right: &Self, right_range: &Range<Key>);
+}
+
+pub trait LazyRangeInitializer<Result: RangeQueryResult<Key>, Key> {
+    fn get(&self, range: &Range<Key>) -> Result;
+}
+
+/// Should be a monoid:
+/// * Identity element: for all op: compose(no_op(), op) == compose(op, no_op()) == op
+/// * Associativity: for all op_1, op_2, op_3: compose(compose(op_1, op_2), op_3) == compose(op_1, compose(op_2, op_3))
+///
+/// Should left act on Result:
+/// * Identity operation: for all r: no_op().apply(r) == r
+/// * Compatibility: for all op_1, op_2, r: op_1.apply(op_2.apply(r)) == compose(op_1, op_2).apply(r)
+pub trait RangeModification<Key> {
+    type Result: RangeQueryResult<Key>;
+
+    fn no_op() -> Self;
+    fn is_no_op(&self) -> bool;
+    fn is_reinitialization(&self) -> bool;
+    fn apply(&self, result: &mut Self::Result, range: &Range<Key>);
+    fn compose(later: &Self, earlier: &mut Self);
+}
+
+pub trait VecReadableVersion<Modification: RangeModification<Key>, Key> {
+    fn get(&self, keys: &Range<Key>) -> Modification::Result;
+}
+
+// TODO: use trait alias when stabilized
+pub trait VecFrozenVersion<Modification: RangeModification<Key>, Key>:
+    Clone + VecReadableVersion<Modification, Key>
+{
+}
+
+impl<
+        T: Clone + VecReadableVersion<Modification, Key>,
+        Modification: RangeModification<Key>,
+        Key,
+    > VecFrozenVersion<Modification, Key> for T
+{
+}
+
+pub trait PersistentVecStorage<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
+    Key,
+>: VecReadableVersion<Modification, Key>
+{
+    fn new(all_keys: Range<Key>, initializer: Initializer) -> Self;
+
+    type FrozenVersion: VecFrozenVersion<Modification, Key>;
+
+    fn modify(&mut self, keys: &Range<Key>, modification: &Modification);
+    fn freeze(&mut self) -> Self::FrozenVersion;
+}

libs/persistent_range_query/src/naive.rs

@@ -0,0 +1,115 @@
+use crate::{
+    LazyRangeInitializer, PersistentVecStorage, RangeModification, RangeQueryResult,
+    VecReadableVersion,
+};
+use std::marker::PhantomData;
+use std::ops::Range;
+use std::rc::Rc;
+
+pub struct NaiveFrozenVersion<Modification: RangeModification<Key>, Key> {
+    all_keys: Range<Key>,
+    values: Rc<Box<Vec<Modification::Result>>>,
+}
+
+pub trait IndexableKey: Clone {
+    fn index(all_keys: &Range<Self>, key: &Self) -> usize;
+    fn element_range(all_keys: &Range<Self>, index: usize) -> Range<Self>;
+}
+
+fn get<Modification: RangeModification<Key>, Key: IndexableKey>(
+    all_keys: &Range<Key>,
+    values: &Vec<Modification::Result>,
+    keys: &Range<Key>,
+) -> Modification::Result {
+    let mut result = Modification::Result::new_for_empty_range();
+    let mut result_range = keys.start.clone()..keys.start.clone();
+    for index in
+        IndexableKey::index(&all_keys, &keys.start)..IndexableKey::index(&all_keys, &keys.end)
+    {
+        let element_range = IndexableKey::element_range(&all_keys, index);
+        Modification::Result::add(&mut result, &result_range, &values[index], &element_range);
+        result_range.end = element_range.end;
+    }
+    result
+}
+
+impl<Modification: RangeModification<Key>, Key: IndexableKey> VecReadableVersion<Modification, Key>
+    for NaiveFrozenVersion<Modification, Key>
+{
+    fn get(&self, keys: &Range<Key>) -> Modification::Result {
+        get::<Modification, Key>(&self.all_keys, &self.values, keys)
+    }
+}
+
+// Manual implementation of `Clone` becase `derive` requires `Modification: Clone`
+impl<Modification: RangeModification<Key>, Key: Clone> Clone
+    for NaiveFrozenVersion<Modification, Key>
+{
+    fn clone(&self) -> Self {
+        Self {
+            all_keys: self.all_keys.clone(),
+            values: self.values.clone(),
+        }
+    }
+}
+
+// TODO: is it at all possible to store previous versions in this struct,
+// without any Rc<>?
+pub struct NaiveVecStorage<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
+    Key: IndexableKey,
+> {
+    all_keys: Range<Key>,
+    last_version: Vec<Modification::Result>,
+    _initializer: PhantomData<Initializer>,
+}
+
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: IndexableKey,
+    > VecReadableVersion<Modification, Key> for NaiveVecStorage<Modification, Initializer, Key>
+{
+    fn get(&self, keys: &Range<Key>) -> Modification::Result {
+        get::<Modification, Key>(&self.all_keys, &self.last_version, keys)
+    }
+}
+
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: IndexableKey,
+    > PersistentVecStorage<Modification, Initializer, Key>
+    for NaiveVecStorage<Modification, Initializer, Key>
+{
+    fn new(all_keys: Range<Key>, initializer: Initializer) -> Self {
+        let mut values = Vec::with_capacity(IndexableKey::index(&all_keys, &all_keys.end));
+        for index in 0..values.capacity() {
+            values.push(initializer.get(&IndexableKey::element_range(&all_keys, index)));
+        }
+        NaiveVecStorage {
+            all_keys,
+            last_version: values,
+            _initializer: PhantomData,
+        }
+    }
+
+    type FrozenVersion = NaiveFrozenVersion<Modification, Key>;
+
+    fn modify(&mut self, keys: &Range<Key>, modification: &Modification) {
+        for index in IndexableKey::index(&self.all_keys, &keys.start)
+            ..IndexableKey::index(&self.all_keys, &keys.end)
+        {
+            let element_range = IndexableKey::element_range(&self.all_keys, index);
+            modification.apply(&mut self.last_version[index], &element_range);
+        }
+    }
+
+    fn freeze(&mut self) -> Self::FrozenVersion {
+        NaiveFrozenVersion::<Modification, Key> {
+            all_keys: self.all_keys.clone(),
+            values: Rc::new(Box::new(self.last_version.clone())),
+        }
+    }
+}

libs/persistent_range_query/src/ops/mod.rs

@@ -0,0 +1,14 @@
+pub mod rsq;
+
+#[derive(Copy, Clone, Debug)]
+pub struct SameElementsInitializer<T> {
+    initial_element_value: T,
+}
+
+impl<T> SameElementsInitializer<T> {
+    pub fn new(initial_element_value: T) -> Self {
+        SameElementsInitializer {
+            initial_element_value,
+        }
+    }
+}

libs/persistent_range_query/src/ops/rsq.rs

@@ -0,0 +1,118 @@
+//! # Range Sum Query
+
+use crate::ops::SameElementsInitializer;
+use crate::{LazyRangeInitializer, RangeModification, RangeQueryResult};
+use std::borrow::Borrow;
+use std::ops::{Add, AddAssign, Range};
+
+// TODO: commutative Add
+
+#[derive(Clone, Copy, Debug)]
+pub struct SumResult<T> {
+    sum: T,
+}
+
+impl<T> SumResult<T> {
+    pub fn sum(&self) -> &T {
+        &self.sum
+    }
+}
+
+impl<T: Clone + for<'a> AddAssign<&'a T> + From<u8>, Key> RangeQueryResult<Key> for SumResult<T>
+where
+    for<'a> &'a T: Add<&'a T, Output = T>,
+{
+    fn new_for_empty_range() -> Self {
+        SumResult { sum: 0.into() }
+    }
+
+    fn combine(
+        left: &Self,
+        _left_range: &Range<Key>,
+        right: &Self,
+        _right_range: &Range<Key>,
+    ) -> Self {
+        SumResult {
+            sum: &left.sum + &right.sum,
+        }
+    }
+
+    fn add(left: &mut Self, _left_range: &Range<Key>, right: &Self, _right_range: &Range<Key>) {
+        left.sum += &right.sum
+    }
+}
+
+pub trait SumOfSameElements<Key> {
+    fn sum(initial_element_value: &Self, keys: &Range<Key>) -> Self;
+}
+
+impl<T: SumOfSameElements<Key>, TB: Borrow<T>, Key> LazyRangeInitializer<SumResult<T>, Key>
+    for SameElementsInitializer<TB>
+where
+    SumResult<T>: RangeQueryResult<Key>,
+{
+    fn get(&self, range: &Range<Key>) -> SumResult<T> {
+        SumResult {
+            sum: SumOfSameElements::sum(self.initial_element_value.borrow(), range),
+        }
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub enum AddAssignModification<T> {
+    None,
+    Add(T),
+    Assign(T),
+}
+
+impl<T: Clone + for<'a> AddAssign<&'a T>, Key> RangeModification<Key> for AddAssignModification<T>
+where
+    SumResult<T>: RangeQueryResult<Key>,
+    for<'a> SameElementsInitializer<&'a T>: LazyRangeInitializer<SumResult<T>, Key>,
+{
+    type Result = SumResult<T>;
+
+    fn no_op() -> Self {
+        AddAssignModification::None
+    }
+
+    fn is_no_op(&self) -> bool {
+        match self {
+            AddAssignModification::None => true,
+            _ => false,
+        }
+    }
+
+    fn is_reinitialization(&self) -> bool {
+        match self {
+            AddAssignModification::Assign(_) => true,
+            _ => false,
+        }
+    }
+
+    fn apply(&self, result: &mut SumResult<T>, range: &Range<Key>) {
+        use AddAssignModification::*;
+        match self {
+            None => {}
+            Add(x) | Assign(x) => {
+                let to_add = SameElementsInitializer::new(x).get(range).sum;
+                if let Assign(_) = self {
+                    result.sum = to_add;
+                } else {
+                    result.sum += &to_add;
+                }
+            }
+        }
+    }
+
+    fn compose(later: &Self, earlier: &mut Self) {
+        use AddAssignModification::*;
+        match (later, earlier) {
+            (_, e @ None) => *e = later.clone(),
+            (None, _) => {}
+            (Assign(_), e) => *e = later.clone(),
+            (Add(x), Add(y)) => *y += x,
+            (Add(x), Assign(value)) => *value += x,
+        }
+    }
+}

libs/persistent_range_query/src/segment_tree.rs

@@ -0,0 +1,255 @@
+//! # Segment Tree
+//! It is a competitive programming folklore data structure. Do not confuse with the interval tree.
+
+use crate::{LazyRangeInitializer, PersistentVecStorage, RangeQueryResult, VecReadableVersion};
+use std::ops::Range;
+use std::rc::Rc;
+
+pub trait MidpointableKey: Clone + Ord + Sized {
+    fn midpoint(range: &Range<Self>) -> Self;
+}
+
+pub trait RangeModification<Key>: Clone + crate::RangeModification<Key> {}
+
+// TODO: use trait alias when stabilized
+impl<T: Clone + crate::RangeModification<Key>, Key> RangeModification<Key> for T {}
+
+#[derive(Debug)]
+struct Node<Modification: RangeModification<Key>, Key> {
+    result: Modification::Result,
+    modify_children: Modification,
+    left: Option<Rc<Self>>,
+    right: Option<Rc<Self>>,
+}
+
+// Manual implementation because we don't need `Key: Clone` for this, unlike with `derive`.
+impl<Modification: RangeModification<Key>, Key> Clone for Node<Modification, Key> {
+    fn clone(&self) -> Self {
+        Node {
+            result: self.result.clone(),
+            modify_children: self.modify_children.clone(),
+            left: self.left.clone(),
+            right: self.right.clone(),
+        }
+    }
+}
+
+impl<Modification: RangeModification<Key>, Key> Node<Modification, Key> {
+    fn new<Initializer: LazyRangeInitializer<Modification::Result, Key>>(
+        range: &Range<Key>,
+        initializer: &Initializer,
+    ) -> Self {
+        Node {
+            result: initializer.get(range),
+            modify_children: Modification::no_op(),
+            left: None,
+            right: None,
+        }
+    }
+
+    pub fn apply(&mut self, modification: &Modification, range: &Range<Key>) {
+        modification.apply(&mut self.result, range);
+        Modification::compose(modification, &mut self.modify_children);
+        if self.modify_children.is_reinitialization() {
+            self.left = None;
+            self.right = None;
+        }
+    }
+
+    pub fn force_children<Initializer: LazyRangeInitializer<Modification::Result, Key>>(
+        &mut self,
+        initializer: &Initializer,
+        range_left: &Range<Key>,
+        range_right: &Range<Key>,
+    ) {
+        let left = Rc::make_mut(
+            self.left
+                .get_or_insert_with(|| Rc::new(Node::new(&range_left, initializer))),
+        );
+        let right = Rc::make_mut(
+            self.right
+                .get_or_insert_with(|| Rc::new(Node::new(&range_right, initializer))),
+        );
+        left.apply(&self.modify_children, &range_left);
+        right.apply(&self.modify_children, &range_right);
+        self.modify_children = Modification::no_op();
+    }
+
+    pub fn recalculate_from_children(&mut self, range_left: &Range<Key>, range_right: &Range<Key>) {
+        assert!(self.modify_children.is_no_op());
+        assert!(self.left.is_some());
+        assert!(self.right.is_some());
+        self.result = Modification::Result::combine(
+            &self.left.as_ref().unwrap().result,
+            &range_left,
+            &self.right.as_ref().unwrap().result,
+            &range_right,
+        );
+    }
+}
+
+fn split_range<Key: MidpointableKey>(range: &Range<Key>) -> (Range<Key>, Range<Key>) {
+    let range_left = range.start.clone()..MidpointableKey::midpoint(range);
+    let range_right = range_left.end.clone()..range.end.clone();
+    (range_left, range_right)
+}
+
+pub struct PersistentSegmentTreeVersion<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
+    Key: Clone,
+> {
+    root: Rc<Node<Modification, Key>>,
+    all_keys: Range<Key>,
+    initializer: Rc<Initializer>,
+}
+
+// Manual implementation because we don't need `Key: Clone` for this, unlike with `derive`.
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: Clone,
+    > Clone for PersistentSegmentTreeVersion<Modification, Initializer, Key>
+{
+    fn clone(&self) -> Self {
+        Self {
+            root: self.root.clone(),
+            all_keys: self.all_keys.clone(),
+            initializer: self.initializer.clone(),
+        }
+    }
+}
+
+fn get<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
+    Key: MidpointableKey,
+>(
+    node: &mut Rc<Node<Modification, Key>>,
+    node_keys: &Range<Key>,
+    initializer: &Initializer,
+    keys: &Range<Key>,
+) -> Modification::Result {
+    if node_keys.end <= keys.start || keys.end <= node_keys.start {
+        return Modification::Result::new_for_empty_range();
+    }
+    if keys.start <= node_keys.start && node_keys.end <= keys.end {
+        return node.result.clone();
+    }
+    let node = Rc::make_mut(node);
+    let (left_keys, right_keys) = split_range(node_keys);
+    node.force_children(initializer, &left_keys, &right_keys);
+    let mut result = get(node.left.as_mut().unwrap(), &left_keys, initializer, keys);
+    Modification::Result::add(
+        &mut result,
+        &left_keys,
+        &get(node.right.as_mut().unwrap(), &right_keys, initializer, keys),
+        &right_keys,
+    );
+    result
+}
+
+fn modify<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
+    Key: MidpointableKey,
+>(
+    node: &mut Rc<Node<Modification, Key>>,
+    node_keys: &Range<Key>,
+    initializer: &Initializer,
+    keys: &Range<Key>,
+    modification: &Modification,
+) {
+    if modification.is_no_op() || node_keys.end <= keys.start || keys.end <= node_keys.start {
+        return;
+    }
+    let node = Rc::make_mut(node);
+    if keys.start <= node_keys.start && node_keys.end <= keys.end {
+        node.apply(modification, node_keys);
+        return;
+    }
+    let (left_keys, right_keys) = split_range(node_keys);
+    node.force_children(initializer, &left_keys, &right_keys);
+    modify(
+        node.left.as_mut().unwrap(),
+        &left_keys,
+        initializer,
+        keys,
+        &modification,
+    );
+    modify(
+        node.right.as_mut().unwrap(),
+        &right_keys,
+        initializer,
+        keys,
+        &modification,
+    );
+    node.recalculate_from_children(&left_keys, &right_keys);
+}
+
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: MidpointableKey,
+    > VecReadableVersion<Modification, Key>
+    for PersistentSegmentTreeVersion<Modification, Initializer, Key>
+{
+    fn get(&self, keys: &Range<Key>) -> Modification::Result {
+        get(
+            &mut self.root.clone(), // TODO: do not always force a branch
+            &self.all_keys,
+            self.initializer.as_ref(),
+            keys,
+        )
+    }
+}
+
+pub struct PersistentSegmentTree<
+    Modification: RangeModification<Key>,
+    Initializer: LazyRangeInitializer<Modification::Result, Key>,
+    Key: MidpointableKey,
+>(PersistentSegmentTreeVersion<Modification, Initializer, Key>);
+
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: MidpointableKey,
+    > VecReadableVersion<Modification, Key>
+    for PersistentSegmentTree<Modification, Initializer, Key>
+{
+    fn get(&self, keys: &Range<Key>) -> Modification::Result {
+        self.0.get(keys)
+    }
+}
+
+impl<
+        Modification: RangeModification<Key>,
+        Initializer: LazyRangeInitializer<Modification::Result, Key>,
+        Key: MidpointableKey,
+    > PersistentVecStorage<Modification, Initializer, Key>
+    for PersistentSegmentTree<Modification, Initializer, Key>
+{
+    fn new(all_keys: Range<Key>, initializer: Initializer) -> Self {
+        PersistentSegmentTree(PersistentSegmentTreeVersion {
+            root: Rc::new(Node::new(&all_keys, &initializer)),
+            all_keys: all_keys,
+            initializer: Rc::new(initializer),
+        })
+    }
+
+    type FrozenVersion = PersistentSegmentTreeVersion<Modification, Initializer, Key>;
+
+    fn modify(&mut self, keys: &Range<Key>, modification: &Modification) {
+        modify(
+            &mut self.0.root, // TODO: do not always force a branch
+            &self.0.all_keys,
+            self.0.initializer.as_ref(),
+            keys,
+            modification,
+        )
+    }
+
+    fn freeze(&mut self) -> Self::FrozenVersion {
+        self.0.clone()
+    }
+}

libs/persistent_range_query/tests/layer_map_test.rs

@@ -0,0 +1,295 @@
+use persistent_range_query::naive::{IndexableKey, NaiveVecStorage};
+use persistent_range_query::ops::SameElementsInitializer;
+use persistent_range_query::segment_tree::{MidpointableKey, PersistentSegmentTree};
+use persistent_range_query::{
+    LazyRangeInitializer, PersistentVecStorage, RangeModification, RangeQueryResult,
+    VecReadableVersion,
+};
+use std::cmp::Ordering;
+use std::ops::Range;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
+struct PageIndex(u32);
+type LayerId = String;
+
+impl IndexableKey for PageIndex {
+    fn index(all_keys: &Range<Self>, key: &Self) -> usize {
+        (key.0 as usize) - (all_keys.start.0 as usize)
+    }
+
+    fn element_range(all_keys: &Range<Self>, index: usize) -> Range<Self> {
+        PageIndex(all_keys.start.0 + index as u32)..PageIndex(all_keys.start.0 + index as u32 + 1)
+    }
+}
+
+impl MidpointableKey for PageIndex {
+    fn midpoint(range: &Range<Self>) -> Self {
+        PageIndex(range.start.0 + (range.end.0 - range.start.0) / 2)
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+struct LayerMapInformation {
+    // Only make sense for a range of length 1.
+    last_layer: Option<LayerId>,
+    last_image_layer: Option<LayerId>,
+    // Work for all ranges
+    max_delta_layers: (usize, Range<PageIndex>),
+}
+
+impl LayerMapInformation {
+    fn last_layers(&self) -> (&Option<LayerId>, &Option<LayerId>) {
+        (&self.last_layer, &self.last_image_layer)
+    }
+
+    fn max_delta_layers(&self) -> &(usize, Range<PageIndex>) {
+        &self.max_delta_layers
+    }
+}
+
+fn merge_ranges(left: &Range<PageIndex>, right: &Range<PageIndex>) -> Range<PageIndex> {
+    if left.is_empty() {
+        right.clone()
+    } else if right.is_empty() {
+        left.clone()
+    } else if left.end == right.start {
+        left.start..right.end
+    } else {
+        left.clone()
+    }
+}
+
+impl RangeQueryResult<PageIndex> for LayerMapInformation {
+    fn new_for_empty_range() -> Self {
+        LayerMapInformation {
+            last_layer: None,
+            last_image_layer: None,
+            max_delta_layers: (0, PageIndex(0)..PageIndex(0)),
+        }
+    }
+
+    fn combine(
+        left: &Self,
+        _left_range: &Range<PageIndex>,
+        right: &Self,
+        _right_range: &Range<PageIndex>,
+    ) -> Self {
+        // Note that either range may be empty.
+        LayerMapInformation {
+            last_layer: left
+                .last_layer
+                .as_ref()
+                .or_else(|| right.last_layer.as_ref())
+                .cloned(),
+            last_image_layer: left
+                .last_image_layer
+                .as_ref()
+                .or_else(|| right.last_image_layer.as_ref())
+                .cloned(),
+            max_delta_layers: match left.max_delta_layers.0.cmp(&right.max_delta_layers.0) {
+                Ordering::Less => right.max_delta_layers.clone(),
+                Ordering::Greater => left.max_delta_layers.clone(),
+                Ordering::Equal => (
+                    left.max_delta_layers.0,
+                    merge_ranges(&left.max_delta_layers.1, &right.max_delta_layers.1),
+                ),
+            },
+        }
+    }
+
+    fn add(
+        left: &mut Self,
+        left_range: &Range<PageIndex>,
+        right: &Self,
+        right_range: &Range<PageIndex>,
+    ) {
+        *left = Self::combine(&left, left_range, right, right_range);
+    }
+}
+
+#[derive(Clone, Debug)]
+struct AddDeltaLayers {
+    last_layer: LayerId,
+    count: usize,
+}
+
+#[derive(Clone, Debug)]
+struct LayerMapModification {
+    add_image_layer: Option<LayerId>,
+    add_delta_layers: Option<AddDeltaLayers>,
+}
+
+impl LayerMapModification {
+    fn add_image_layer(layer: impl Into<LayerId>) -> Self {
+        LayerMapModification {
+            add_image_layer: Some(layer.into()),
+            add_delta_layers: None,
+        }
+    }
+
+    fn add_delta_layer(layer: impl Into<LayerId>) -> Self {
+        LayerMapModification {
+            add_image_layer: None,
+            add_delta_layers: Some(AddDeltaLayers {
+                last_layer: layer.into(),
+                count: 1,
+            }),
+        }
+    }
+}
+
+impl RangeModification<PageIndex> for LayerMapModification {
+    type Result = LayerMapInformation;
+
+    fn no_op() -> Self {
+        LayerMapModification {
+            add_image_layer: None,
+            add_delta_layers: None,
+        }
+    }
+
+    fn is_no_op(&self) -> bool {
+        self.add_image_layer.is_none() && self.add_delta_layers.is_none()
+    }
+
+    fn is_reinitialization(&self) -> bool {
+        self.add_image_layer.is_some()
+    }
+
+    fn apply(&self, result: &mut Self::Result, range: &Range<PageIndex>) {
+        if let Some(layer) = &self.add_image_layer {
+            result.last_layer = Some(layer.clone());
+            result.last_image_layer = Some(layer.clone());
+            result.max_delta_layers = (0, range.clone());
+        }
+        if let Some(AddDeltaLayers { last_layer, count }) = &self.add_delta_layers {
+            result.last_layer = Some(last_layer.clone());
+            result.max_delta_layers.0 += count;
+        }
+    }
+
+    fn compose(later: &Self, earlier: &mut Self) {
+        if later.add_image_layer.is_some() {
+            *earlier = later.clone();
+            return;
+        }
+        if let Some(AddDeltaLayers { last_layer, count }) = &later.add_delta_layers {
+            let res = earlier.add_delta_layers.get_or_insert(AddDeltaLayers {
+                last_layer: LayerId::default(),
+                count: 0,
+            });
+            res.last_layer = last_layer.clone();
+            res.count += count;
+        }
+    }
+}
+
+impl LazyRangeInitializer<LayerMapInformation, PageIndex> for SameElementsInitializer<()> {
+    fn get(&self, range: &Range<PageIndex>) -> LayerMapInformation {
+        LayerMapInformation {
+            last_layer: None,
+            last_image_layer: None,
+            max_delta_layers: (0, range.clone()),
+        }
+    }
+}
+
+fn test_layer_map<
+    S: PersistentVecStorage<LayerMapModification, SameElementsInitializer<()>, PageIndex>,
+>() {
+    let mut s = S::new(
+        PageIndex(0)..PageIndex(100),
+        SameElementsInitializer::new(()),
+    );
+    s.modify(
+        &(PageIndex(0)..PageIndex(70)),
+        &LayerMapModification::add_image_layer("Img0..70"),
+    );
+    s.modify(
+        &(PageIndex(50)..PageIndex(100)),
+        &LayerMapModification::add_image_layer("Img50..100"),
+    );
+    s.modify(
+        &(PageIndex(10)..PageIndex(60)),
+        &LayerMapModification::add_delta_layer("Delta10..60"),
+    );
+    let s_before_last_delta = s.freeze();
+    s.modify(
+        &(PageIndex(20)..PageIndex(80)),
+        &LayerMapModification::add_delta_layer("Delta20..80"),
+    );
+
+    assert_eq!(
+        s.get(&(PageIndex(5)..PageIndex(6))).last_layers(),
+        (&Some("Img0..70".to_owned()), &Some("Img0..70".to_owned()))
+    );
+    assert_eq!(
+        s.get(&(PageIndex(15)..PageIndex(16))).last_layers(),
+        (
+            &Some("Delta10..60".to_owned()),
+            &Some("Img0..70".to_owned())
+        )
+    );
+    assert_eq!(
+        s.get(&(PageIndex(25)..PageIndex(26))).last_layers(),
+        (
+            &Some("Delta20..80".to_owned()),
+            &Some("Img0..70".to_owned())
+        )
+    );
+    assert_eq!(
+        s.get(&(PageIndex(65)..PageIndex(66))).last_layers(),
+        (
+            &Some("Delta20..80".to_owned()),
+            &Some("Img50..100".to_owned())
+        )
+    );
+    assert_eq!(
+        s.get(&(PageIndex(95)..PageIndex(96))).last_layers(),
+        (
+            &Some("Img50..100".to_owned()),
+            &Some("Img50..100".to_owned())
+        )
+    );
+
+    assert_eq!(
+        s.get(&(PageIndex(0)..PageIndex(100))).max_delta_layers(),
+        &(2, PageIndex(20)..PageIndex(60)),
+    );
+    assert_eq!(
+        *s_before_last_delta
+            .get(&(PageIndex(0)..PageIndex(100)))
+            .max_delta_layers(),
+        (1, PageIndex(10)..PageIndex(60)),
+    );
+
+    assert_eq!(
+        *s.get(&(PageIndex(10)..PageIndex(30))).max_delta_layers(),
+        (2, PageIndex(20)..PageIndex(30))
+    );
+    assert_eq!(
+        *s.get(&(PageIndex(10)..PageIndex(20))).max_delta_layers(),
+        (1, PageIndex(10)..PageIndex(20))
+    );
+
+    assert_eq!(
+        *s.get(&(PageIndex(70)..PageIndex(80))).max_delta_layers(),
+        (1, PageIndex(70)..PageIndex(80))
+    );
+    assert_eq!(
+        *s_before_last_delta
+            .get(&(PageIndex(70)..PageIndex(80)))
+            .max_delta_layers(),
+        (0, PageIndex(70)..PageIndex(80))
+    );
+}
+
+#[test]
+fn test_naive() {
+    test_layer_map::<NaiveVecStorage<_, _, _>>();
+}
+
+#[test]
+fn test_segment_tree() {
+    test_layer_map::<PersistentSegmentTree<_, _, _>>();
+}

libs/persistent_range_query/tests/rsq_test.rs

@@ -0,0 +1,116 @@
+use persistent_range_query::naive::*;
+use persistent_range_query::ops::rsq::AddAssignModification::Add;
+use persistent_range_query::ops::rsq::*;
+use persistent_range_query::ops::SameElementsInitializer;
+use persistent_range_query::segment_tree::{MidpointableKey, PersistentSegmentTree};
+use persistent_range_query::{PersistentVecStorage, VecReadableVersion};
+use rand::{Rng, SeedableRng};
+use std::ops::Range;
+
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
+struct K(u16);
+
+impl IndexableKey for K {
+    fn index(all_keys: &Range<Self>, key: &Self) -> usize {
+        (key.0 as usize) - (all_keys.start.0 as usize)
+    }
+
+    fn element_range(all_keys: &Range<Self>, index: usize) -> Range<Self> {
+        K(all_keys.start.0 + index as u16)..K(all_keys.start.0 + index as u16 + 1)
+    }
+}
+
+impl SumOfSameElements<K> for i32 {
+    fn sum(initial_element_value: &Self, keys: &Range<K>) -> Self {
+        initial_element_value * (keys.end.0 - keys.start.0) as Self
+    }
+}
+
+impl MidpointableKey for K {
+    fn midpoint(range: &Range<Self>) -> Self {
+        K(range.start.0 + (range.end.0 - range.start.0) / 2)
+    }
+}
+
+fn test_storage<
+    S: PersistentVecStorage<AddAssignModification<i32>, SameElementsInitializer<i32>, K>,
+>() {
+    let mut s = S::new(K(0)..K(12), SameElementsInitializer::new(0i32));
+    assert_eq!(*s.get(&(K(0)..K(12))).sum(), 0);
+
+    s.modify(&(K(2)..K(5)), &AddAssignModification::Add(3));
+    assert_eq!(*s.get(&(K(0)..K(12))).sum(), 3 + 3 + 3);
+    let s_old = s.freeze();
+
+    s.modify(&(K(3)..K(6)), &AddAssignModification::Assign(10));
+    assert_eq!(*s.get(&(K(0)..K(12))).sum(), 3 + 10 + 10 + 10);
+
+    s.modify(&(K(4)..K(7)), &AddAssignModification::Add(2));
+    assert_eq!(*s.get(&(K(0)..K(12))).sum(), 3 + 10 + 12 + 12 + 2);
+
+    assert_eq!(*s.get(&(K(4)..K(6))).sum(), 12 + 12);
+    assert_eq!(*s_old.get(&(K(4)..K(6))).sum(), 3);
+}
+
+#[test]
+fn test_naive() {
+    test_storage::<NaiveVecStorage<_, _, _>>();
+}
+
+#[test]
+fn test_segment_tree() {
+    test_storage::<PersistentSegmentTree<_, _, _>>();
+}
+
+#[test]
+fn test_stress() {
+    const LEN: u16 = 17_238;
+    const OPERATIONS: i32 = 20_000;
+
+    let mut rng = rand::rngs::StdRng::seed_from_u64(0);
+    let mut naive: NaiveVecStorage<AddAssignModification<i32>, _, _> =
+        NaiveVecStorage::new(K(0)..K(LEN), SameElementsInitializer::new(2i32));
+    let mut segm_tree: PersistentSegmentTree<AddAssignModification<i32>, _, _> =
+        PersistentSegmentTree::new(K(0)..K(LEN), SameElementsInitializer::new(2i32));
+
+    fn gen_range(rng: &mut impl Rng) -> Range<K> {
+        let l: u16 = rng.gen_range(0..LEN);
+        let r: u16 = rng.gen_range(0..LEN);
+        if l <= r {
+            K(l)..K(r)
+        } else {
+            K(r)..K(l)
+        }
+    }
+
+    for _ in 0..2 {
+        let checksum_range = gen_range(&mut rng);
+        let checksum_before: i32 = *naive.get(&checksum_range).sum();
+        assert_eq!(checksum_before, *segm_tree.get(&checksum_range).sum());
+
+        let naive_before = naive.freeze();
+        let segm_tree_before = segm_tree.freeze();
+        assert_eq!(checksum_before, *naive_before.get(&checksum_range).sum());
+        assert_eq!(checksum_before, *segm_tree.get(&checksum_range).sum());
+
+        for _ in 0..OPERATIONS {
+            {
+                let range = gen_range(&mut rng);
+                assert_eq!(naive.get(&range).sum(), segm_tree.get(&range).sum());
+            }
+            {
+                let range = gen_range(&mut rng);
+                let val = rng.gen_range(-10i32..=10i32);
+                let op = Add(val);
+                naive.modify(&range, &op);
+                segm_tree.modify(&range, &op);
+            }
+        }
+
+        assert_eq!(checksum_before, *naive_before.get(&checksum_range).sum());
+        assert_eq!(
+            checksum_before,
+            *segm_tree_before.get(&checksum_range).sum()
+        );
+    }
+}

pageserver/Cargo.toml

@@ -59,6 +59,7 @@ tracing = "0.1.36"
 url = "2"
 walkdir = "2.3.2"
 
+persistent_range_query = { path = "../libs/persistent_range_query" }
 etcd_broker = { path = "../libs/etcd_broker" }
 metrics = { path = "../libs/metrics" }
 pageserver_api = { path = "../libs/pageserver_api" }
@@ -70,7 +71,6 @@ tenant_size_model = { path = "../libs/tenant_size_model" }
 utils = { path = "../libs/utils" }
 workspace_hack = { version = "0.1", path = "../workspace_hack" }
 rpds = "0.12.0"
-im = "15.1.0"
 
 [dev-dependencies]
 criterion = "0.4"

pageserver/benches/bench_layer_map.rs

@@ -1,9 +1,10 @@
 use anyhow::Result;
 use num_traits::ToPrimitive;
 use pageserver::repository::{Key, Value};
-use pageserver::tenant::bst_layer_map::RetroactiveLayerMap;
+use pageserver::tenant::bst_layer_map::BSTLM;
 use pageserver::tenant::filename::{DeltaFileName, ImageFileName};
 use pageserver::tenant::layer_map::LayerMap;
+use pageserver::tenant::segment_tree_layer_map::STLM;
 use pageserver::tenant::storage_layer::Layer;
 use pageserver::tenant::storage_layer::ValueReconstructResult;
 use pageserver::tenant::storage_layer::ValueReconstructState;
@@ -180,7 +181,6 @@ fn build_layer_map(filename_dump: PathBuf) -> LayerMap {
             panic!("unexpected filename {fname}");
         }
     }
-    layer_map.rebuild_index();
 
     println!("min: {min_lsn}, max: {max_lsn}");
 
@@ -254,18 +254,35 @@ fn bench_from_real_project(c: &mut Criterion) {
 
     // Init bst layer map with the same layers
     let now = Instant::now();
-    let mut bstlm = RetroactiveLayerMap::new();
-    for layer in layer_map.iter_historic_layers() {
-        let kr = layer.get_key_range();
-        let lr = layer.get_lsn_range();
+    let mut bstlm = BSTLM::new();
+    let mut sorted_layers: Vec<_> = layer_map.iter_historic_layers().collect();
+    sorted_layers.sort_by(|a, b| {
+        a.get_lsn_range().start.cmp(&b.get_lsn_range().start)
+    });
+    for layer in sorted_layers {
+        if layer.is_incremental() {
+            // TODO check if they're sorted
+            let kr = layer.get_key_range();
+            let lr = layer.get_lsn_range();
 
-        bstlm.insert(
-            kr.start.to_i128()..kr.end.to_i128(),
-            lr.start.0..lr.end.0,
-            format!("Layer {}", lr.start.0),
-        );
+            bstlm.insert(
+                kr.start.to_i128(),
+                kr.end.to_i128(),
+                lr.start.0,
+                format!("Layer {}", lr.start.0),
+            );
+        } else {
+            let kr = layer.get_key_range();
+            let lr = layer.get_lsn_range();
+
+            bstlm.insert(
+                kr.start.to_i128(),
+                kr.end.to_i128(),
+                lr.start.0,
+                format!("Layer {}", lr.start.0),
+            );
+        }
     }
-    bstlm.rebuild();
     println!("Finished bst init in {:?}", now.elapsed());
 
     // Choose uniformly distributed queries
@@ -283,7 +300,7 @@ fn bench_from_real_project(c: &mut Criterion) {
     group.bench_function("persistent_bst", |b| {
         b.iter(|| {
             for q in queries.clone().into_iter() {
-                bstlm.query(q.0.to_i128(), q.1 .0);
+                bstlm.query(q.0.to_i128(), q.1.0);
             }
         });
     });
@@ -292,6 +309,7 @@ fn bench_from_real_project(c: &mut Criterion) {
 
 // Benchmark using synthetic data. Arrange image layers on stacked diagonal lines.
 fn bench_sequential(c: &mut Criterion) {
+
     // Init layer map. Create 100_000 layers arranged in 1000 diagonal lines.
     //
     // TODO This code is pretty slow and runs even if we're only running other
@@ -309,23 +327,26 @@ fn bench_sequential(c: &mut Criterion) {
         };
         layer_map.insert_historic(Arc::new(layer));
     }
-    layer_map.rebuild_index();
     println!("Finished layer map init in {:?}", now.elapsed());
 
     // Init bst layer map with the same layers
     let now = Instant::now();
-    let mut bstlm = RetroactiveLayerMap::new();
+    let mut bstlm = BSTLM::new();
     for layer in layer_map.iter_historic_layers() {
-        let kr = layer.get_key_range();
-        let lr = layer.get_lsn_range();
+        if layer.is_incremental() {
+            panic!("AAA");
+        } else {
+            let kr = layer.get_key_range();
+            let lr = layer.get_lsn_range();
 
-        bstlm.insert(
-            kr.start.to_i128()..kr.end.to_i128(),
-            lr.start.0..lr.end.0,
-            format!("Layer {}", lr.start.0),
-        );
+            bstlm.insert(
+                kr.start.to_i128(),
+                kr.end.to_i128(),
+                lr.start.0,
+                format!("Layer {}", lr.start.0),
+            );
+        }
     }
-    bstlm.rebuild();
     println!("Finished bst init in {:?}", now.elapsed());
 
     // Choose 100 uniformly random queries
@@ -347,14 +368,14 @@ fn bench_sequential(c: &mut Criterion) {
     group.bench_function("persistent_bst", |b| {
         b.iter(|| {
             for q in queries.clone().into_iter() {
-                bstlm.query(q.0.to_i128(), q.1 .0);
+                bstlm.query(q.0.to_i128(), q.1.0);
             }
         });
     });
     group.finish();
 }
 
-criterion_group!(group_1, bench_from_captest_env);
-criterion_group!(group_2, bench_from_real_project);
-criterion_group!(group_3, bench_sequential);
-criterion_main!(group_1, group_2, group_3);
+// HACK TODO bring back all the bench functions. I remove
+// them here to avoid initializing.
+criterion_group!(group, bench_from_real_project);
+criterion_main!(group);

pageserver/src/tenant.rs

@@ -73,7 +73,6 @@ use utils::{
 
 mod blob_io;
 pub mod block_io;
-pub mod bst_layer_map;
 mod delta_layer;
 mod disk_btree;
 pub(crate) mod ephemeral_file;
@@ -81,6 +80,8 @@ pub mod filename;
 mod image_layer;
 mod inmemory_layer;
 pub mod layer_map;
+pub mod bst_layer_map;
+pub mod segment_tree_layer_map;
 
 pub mod metadata;
 mod par_fsync;

pageserver/src/tenant/bst_layer_map.rs

@@ -1,407 +1,111 @@
 use std::collections::BTreeMap;
-use std::ops::Range;
 
-// TODO drop rpds. So far `im` looks 30% faster.
-use rpds::RedBlackTreeMapSync;
-use im::OrdMap;
+// TODO the `im` crate has 20x more downloads and also has persistent
+// BTree. See if it's better. (What's "fully persistent?")
+use rpds::RedBlackTreeMap;
 
-/// Layer map implemented using persistent/immutable binary search tree.
-/// It supports historical queries, but no retroactive inserts. For that
-/// see RetroactiveLayerMap.
-///
-/// Layer type is abstracted as Value to make unit testing easier.
-pub struct PersistentLayerMap<Value> {
-    /// Mapping key to the latest layer (if any) until the next key.
-    /// We use the Sync version of the map because we want Self to
-    /// be Sync.
-    ///
-    /// TODO Separate Head into its own struct LatestLayerMap
-    /// TODO Merge historic with retroactive, into HistoricLayerMap
-    /// TODO Maintain a pair of heads, one for images, one for deltas.
-    ///      This way we can query both of them with one BTreeMap query.
-    head: OrdMap<i128, Option<(u64, Value)>>,
+
+/// Layer map implemented using persistent binary search tree.
+/// This implementation is only good enough to run benchmarks,
+/// so it's missing unnecessary details. Values are String for now.
+pub struct BSTLM {
+    /// Mapping key to the latest layer (if any) until the next key
+    head: RedBlackTreeMap<i128, Option<String>>,
 
     /// All previous states of `self.head`
-    ///
-    /// TODO: Sorted Vec + binary search could be slightly faster.
-    historic: BTreeMap<u64, OrdMap<i128, Option<(u64, Value)>>>,
+    historic: BTreeMap<u64, RedBlackTreeMap<i128, Option<String>>>,
 }
 
-impl<Value: std::fmt::Debug> std::fmt::Debug for PersistentLayerMap<Value> {
+impl std::fmt::Debug for BSTLM {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let head_vec: Vec<_> = self.head.iter().collect();
-        write!(f, "PersistentLayerMap: head: {:?}", head_vec)
+        write!(f, "BSTLM: head: {:?}", head_vec)
     }
 }
 
-impl<T: Clone> Default for PersistentLayerMap<T> {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl<Value: Clone> PersistentLayerMap<Value> {
+impl BSTLM {
     pub fn new() -> Self {
-        Self {
-            head: OrdMap::default(),
+        BSTLM {
+            head: RedBlackTreeMap::default(),
             historic: BTreeMap::default(),
         }
     }
 
-    /// Helper function to subdivide the key range without changing any values
-    fn add_node(self: &mut Self, key: i128) {
-        let value = match self.head.range(0..=key).last() {
-            Some((_, Some(v))) => Some(v.clone()),
-            Some((_, None)) => None,
-            None => None,
-        };
-        self.head.insert(key, value);
-    }
+    pub fn insert(self: &mut Self, key_begin: i128, key_end: i128, lsn: u64, value: String) {
+        // TODO check for off-by-one errors
 
-    pub fn insert(self: &mut Self, key: Range<i128>, lsn: Range<u64>, value: Value) {
-        // It's only a persistent map, not a retroactive one
+        // It's a persistent map, not a retroactive one
         if let Some(last_entry) = self.historic.iter().rev().next() {
             let last_lsn = last_entry.0;
-            if lsn.start == *last_lsn {
+            if lsn == *last_lsn {
                 // TODO there are edge cases to take care of
             }
-            if lsn.start < *last_lsn {
-                panic!("unexpected retroactive insert");
+            if lsn < *last_lsn {
+                todo!("smaller lsn not implemented yet")
             }
         }
 
         // NOTE The order of the following lines is important!!
 
-        // Add nodes at endpoints
-        self.add_node(key.start);
-        self.add_node(key.end);
+        // Preserve information after right endpoint
+        let value_at_end = match self.head.range(0..key_end).last() {
+            Some((_, Some(v))) => Some(v.clone()),
+            Some((_, None)) => None,
+            None => None,
+        };
+        self.head.insert_mut(key_end, value_at_end);
 
-        // Raise the height where necessary
-        //
-        // NOTE This loop is worst case O(N), but amortized O(log N) in the special
-        // case when rectangles have no height. In practice I don't think we'll see
-        // the kind of layer intersections needed to trigger O(N) behavior. If we
-        // do it can be fixed using lazy propagation.
-        let mut to_update = Vec::new();
-        let mut to_remove = Vec::new();
-        let mut prev_covered = false;
-        for (k, node) in self.head.range(key.clone()) {
-            let needs_cover = match node {
-                None => true,
-                Some((h, _)) => h < &lsn.end,
-            };
-            if needs_cover {
-                match prev_covered {
-                    true => to_remove.push(k.clone()),
-                    false => to_update.push(k.clone()),
-                }
-            }
-            prev_covered = needs_cover;
-        }
-        if !prev_covered {
-            to_remove.push(key.end);
-        }
-        for k in to_update {
-            self.head
-                .insert(k.clone(), Some((lsn.end.clone(), value.clone())));
-        }
-        for k in to_remove {
-            self.head.remove(&k);
+        // Insert the left endpoint
+        self.head.insert_mut(key_begin, Some(value.clone()));
+
+        // Cover the inside of the interval
+        let to_remove: Vec<_> = self.head.range((key_begin + 1)..key_end)
+            .map(|(k, _)| k.clone())
+            .collect();
+        for key in to_remove {
+            self.head.remove_mut(&key);
         }
 
         // Remember history. Clone is O(1)
-        self.historic.insert(lsn.start, self.head.clone());
+        self.historic.insert(lsn, self.head.clone());
     }
 
-    pub fn query(self: &Self, key: i128, lsn: u64) -> Option<Value> {
+    pub fn query(self: &Self, key: i128, lsn: u64) -> Option<&String> {
+        // TODO check for off-by-one errors
+
         let version = self.historic.range(0..=lsn).rev().next()?.1;
-        version
-            .get_prev(&key)?
-            // .range(0..=key).rev().next()?
-            // NOTE The canonical way to do this in other crates is
-            //      `.range(0..=key).rev.next()` and `im` supports this
-            //      API but it's 2x slower than `.get_prev(&key)`.
-            .1
-            .as_ref()
-            .map(|(_, v)| v.clone())
+        version.range(0..=key).rev().next()?.1.as_ref()
     }
 
-    pub fn trim(self: &mut Self, begin: &u64) {
-        self.historic.split_off(begin);
-        self.head = self
-            .historic
-            .iter()
-            .rev()
-            .next()
-            .map(|(_, v)| v.clone())
-            .unwrap_or_default();
-    }
+    // TODO Add API for delta layers with lsn range.
+    //      The easy solution is to only store images, and then from every
+    //      image point to deltas on top of it. There might be something
+    //      nicer but we have this solution as backup.
 }
 
-/// Basic test for the immutable bst library, just to show usage.
+
 #[test]
-fn test_immutable_bst_dependency() {
-    let map = RedBlackTreeMapSync::<i32, i32>::default();
-
-    let mut v1 = map.clone();
-    let v2 = map.insert(1, 5);
-
-    // We can query current and past versions of key 1
-    assert_eq!(v1.get(&1), None);
-    assert_eq!(v2.get(&1), Some(&5));
-
-    // We can mutate old state, but it creates a branch.
-    // It doesn't retroactively change future versions.
-    v1.insert_mut(2, 6);
-    assert_eq!(v1.get(&2), Some(&6));
-    assert_eq!(v2.get(&2), None);
-}
-
-/// This is the most basic test that demonstrates intended usage.
-/// All layers in this test have height 1.
-#[test]
-fn test_persistent_simple() {
-    let mut map = PersistentLayerMap::<String>::new();
-    map.insert(0..5, 100..101, "Layer 1".to_string());
-    map.insert(3..9, 110..111, "Layer 2".to_string());
-    map.insert(5..6, 120..121, "Layer 3".to_string());
+fn test_bstlm() {
+    let mut bstlm = BSTLM::new();
+    bstlm.insert(0, 5, 100, "Layer 1".to_string());
+    dbg!(&bstlm);
+    bstlm.insert(3, 9, 110, "Layer 2".to_string());
+    dbg!(&bstlm);
+    bstlm.insert(5, 6, 120, "Layer 3".to_string());
+    dbg!(&bstlm);
 
     // After Layer 1 insertion
-    assert_eq!(map.query(1, 105), Some("Layer 1".to_string()));
-    assert_eq!(map.query(4, 105), Some("Layer 1".to_string()));
+    assert_eq!(bstlm.query(1, 105), Some(&"Layer 1".to_string()));
+    assert_eq!(bstlm.query(4, 105), Some(&"Layer 1".to_string()));
 
     // After Layer 2 insertion
-    assert_eq!(map.query(4, 115), Some("Layer 2".to_string()));
-    assert_eq!(map.query(8, 115), Some("Layer 2".to_string()));
-    assert_eq!(map.query(11, 115), None);
+    assert_eq!(bstlm.query(4, 115), Some(&"Layer 2".to_string()));
+    assert_eq!(bstlm.query(8, 115), Some(&"Layer 2".to_string()));
+    assert_eq!(bstlm.query(11, 115), None);
 
     // After Layer 3 insertion
-    assert_eq!(map.query(4, 125), Some("Layer 2".to_string()));
-    assert_eq!(map.query(5, 125), Some("Layer 3".to_string()));
-    assert_eq!(map.query(7, 125), Some("Layer 2".to_string()));
-}
-
-/// Cover simple off-by-one edge cases
-#[test]
-fn test_off_by_one() {
-    let mut map = PersistentLayerMap::<String>::new();
-    map.insert(3..5, 100..110, "Layer 1".to_string());
-
-    // Check different LSNs
-    assert_eq!(map.query(4, 99), None);
-    assert_eq!(map.query(4, 100), Some("Layer 1".to_string()));
-
-    // Check different keys
-    assert_eq!(map.query(2, 105), None);
-    assert_eq!(map.query(3, 105), Some("Layer 1".to_string()));
-    assert_eq!(map.query(4, 105), Some("Layer 1".to_string()));
-    assert_eq!(map.query(5, 105), None);
-}
-
-/// Cover edge cases where layers begin or end on the same key
-#[test]
-fn test_key_collision() {
-    let mut map = PersistentLayerMap::<String>::new();
-
-    map.insert(3..5, 100..110, "Layer 10".to_string());
-    map.insert(5..8, 100..110, "Layer 11".to_string());
-
-    map.insert(3..4, 200..210, "Layer 20".to_string());
-
-    // Check after layer 11
-    assert_eq!(map.query(2, 105), None);
-    assert_eq!(map.query(3, 105), Some("Layer 10".to_string()));
-    assert_eq!(map.query(5, 105), Some("Layer 11".to_string()));
-    assert_eq!(map.query(7, 105), Some("Layer 11".to_string()));
-    assert_eq!(map.query(8, 105), None);
-
-    // Check after layer 20
-    assert_eq!(map.query(2, 205), None);
-    assert_eq!(map.query(3, 205), Some("Layer 20".to_string()));
-    assert_eq!(map.query(5, 205), Some("Layer 11".to_string()));
-    assert_eq!(map.query(7, 205), Some("Layer 11".to_string()));
-    assert_eq!(map.query(8, 205), None);
-}
-
-/// Test when rectangles have nontrivial height and possibly overlap
-#[test]
-fn test_persistent_overlapping() {
-    let mut map = PersistentLayerMap::<String>::new();
-
-    // Add 3 key-disjoint layers with varying LSN ranges
-    map.insert(1..2, 100..200, "Layer 1".to_string());
-    map.insert(4..5, 110..200, "Layer 2".to_string());
-    map.insert(7..8, 120..300, "Layer 3".to_string());
-
-    // Add wide and short layer
-    map.insert(0..9, 130..199, "Layer 4".to_string());
-
-    // Add wide layer taller than some
-    map.insert(0..9, 140..201, "Layer 5".to_string());
-
-    // Add wide layer taller than all
-    map.insert(0..9, 150..301, "Layer 6".to_string());
-
-    // After layer 4 insertion
-    assert_eq!(map.query(0, 135), Some("Layer 4".to_string()));
-    assert_eq!(map.query(1, 135), Some("Layer 1".to_string()));
-    assert_eq!(map.query(2, 135), Some("Layer 4".to_string()));
-    assert_eq!(map.query(4, 135), Some("Layer 2".to_string()));
-    assert_eq!(map.query(5, 135), Some("Layer 4".to_string()));
-    assert_eq!(map.query(7, 135), Some("Layer 3".to_string()));
-    assert_eq!(map.query(8, 135), Some("Layer 4".to_string()));
-
-    // After layer 5 insertion
-    assert_eq!(map.query(0, 145), Some("Layer 5".to_string()));
-    assert_eq!(map.query(1, 145), Some("Layer 5".to_string()));
-    assert_eq!(map.query(2, 145), Some("Layer 5".to_string()));
-    assert_eq!(map.query(4, 145), Some("Layer 5".to_string()));
-    assert_eq!(map.query(5, 145), Some("Layer 5".to_string()));
-    assert_eq!(map.query(7, 145), Some("Layer 3".to_string()));
-    assert_eq!(map.query(8, 145), Some("Layer 5".to_string()));
-
-    // After layer 6 insertion
-    assert_eq!(map.query(0, 155), Some("Layer 6".to_string()));
-    assert_eq!(map.query(1, 155), Some("Layer 6".to_string()));
-    assert_eq!(map.query(2, 155), Some("Layer 6".to_string()));
-    assert_eq!(map.query(4, 155), Some("Layer 6".to_string()));
-    assert_eq!(map.query(5, 155), Some("Layer 6".to_string()));
-    assert_eq!(map.query(7, 155), Some("Layer 6".to_string()));
-    assert_eq!(map.query(8, 155), Some("Layer 6".to_string()));
-}
-
-/// Layer map that supports:
-/// - efficient historical queries
-/// - efficient append only updates
-/// - tombstones and similar methods for non-latest updates
-/// - compaction/rebuilding to remove tombstones
-///
-/// See this for better retroactive techniques we can try
-/// https://www.youtube.com/watch?v=WqCWghETNDc&t=581s
-///
-/// Layer type is abstracted as Value to make unit testing easier.
-pub struct RetroactiveLayerMap<Value> {
-    /// A persistent layer map that we rebuild when we need to retroactively update
-    map: PersistentLayerMap<Value>,
-
-    /// We buffer insertion into the PersistentLayerMap to decrease the number of rebuilds.
-    /// A value of None means we want to delete this item.
-    buffer: BTreeMap<(u64, u64, i128, i128), Option<Value>>,
-
-    /// All current layers. This is not used for search. Only to make rebuilds easier.
-    layers: BTreeMap<(u64, u64, i128, i128), Value>,
-}
-
-impl<Value: std::fmt::Debug> std::fmt::Debug for RetroactiveLayerMap<Value> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "RetroactiveLayerMap: head: {:?}", self.map)
-    }
-}
-
-impl<T: Clone> Default for RetroactiveLayerMap<T> {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl<Value: Clone> RetroactiveLayerMap<Value> {
-    pub fn new() -> Self {
-        Self {
-            map: PersistentLayerMap::<Value>::new(),
-            buffer: BTreeMap::new(),
-            layers: BTreeMap::new(),
-        }
-    }
-
-    pub fn insert(self: &mut Self, key: Range<i128>, lsn: Range<u64>, value: Value) {
-        self.buffer.insert(
-            (lsn.start, lsn.end, key.start, key.end),
-            Some(value.clone()),
-        );
-    }
-
-    pub fn remove(self: &mut Self, key: Range<i128>, lsn: Range<u64>) {
-        self.buffer
-            .insert((lsn.start, lsn.end, key.start, key.end), None);
-    }
-
-    pub fn rebuild(self: &mut Self) {
-        // Find the first LSN that needs to be rebuilt
-        let rebuild_since: u64 = match self.buffer.iter().next() {
-            Some(((lsn_start, _, _, _), _)) => lsn_start.clone(),
-            None => return, // No need to rebuild if buffer is empty
-        };
-
-        // Apply buffered updates to self.layers
-        self.buffer.retain(|rect, layer| {
-            match layer {
-                Some(l) => {
-                    let existing = self.layers.insert(rect.clone(), l.clone());
-                    if existing.is_some() {
-                        panic!("can't overwrite layer");
-                    }
-                }
-                None => {
-                    let existing = self.layers.remove(rect);
-                    if existing.is_none() {
-                        panic!("invalid layer deletion");
-                    }
-                }
-            };
-            false
-        });
-
-        // Rebuild
-        self.map.trim(&rebuild_since);
-        for ((lsn_start, lsn_end, key_start, key_end), layer) in
-            self.layers.range((rebuild_since, 0, 0, 0)..)
-        {
-            self.map
-                .insert(*key_start..*key_end, *lsn_start..*lsn_end, layer.clone());
-        }
-    }
-
-    pub fn clear(self: &mut Self) {
-        self.map.trim(&0);
-    }
-
-    pub fn query(self: &Self, key: i128, lsn: u64) -> Option<Value> {
-        if !self.buffer.is_empty() {
-            panic!("rebuild pls")
-        }
-
-        self.map.query(key, lsn)
-    }
-}
-
-#[test]
-fn test_retroactive_simple() {
-    let mut map = RetroactiveLayerMap::new();
-
-    // Append some images in increasing LSN order
-    map.insert(0..5, 100..101, "Image 1".to_string());
-    map.insert(3..9, 110..111, "Image 2".to_string());
-    map.insert(4..6, 120..121, "Image 3".to_string());
-    map.insert(8..9, 120..121, "Image 4".to_string());
-
-    // Add a delta layer out of order
-    map.insert(2..5, 105..106, "Delta 1".to_string());
-
-    // Rebuild so we can start querying
-    map.rebuild();
-
-    // Query key 4
-    assert_eq!(map.query(4, 90), None);
-    assert_eq!(map.query(4, 102), Some("Image 1".to_string()));
-    assert_eq!(map.query(4, 107), Some("Delta 1".to_string()));
-    assert_eq!(map.query(4, 115), Some("Image 2".to_string()));
-    assert_eq!(map.query(4, 125), Some("Image 3".to_string()));
-
-    // Remove Image 3
-    map.remove(4..6, 120..121);
-    map.rebuild();
-
-    // Check deletion worked
-    assert_eq!(map.query(4, 125), Some("Image 2".to_string()));
-    assert_eq!(map.query(8, 125), Some("Image 4".to_string()));
+    assert_eq!(bstlm.query(4, 125), Some(&"Layer 2".to_string()));
+    assert_eq!(bstlm.query(5, 125), Some(&"Layer 3".to_string()));
+
+    assert_eq!(bstlm.query(7, 125), Some(&"Layer 2".to_string()));
 }

pageserver/src/tenant/layer_map.rs

@@ -28,8 +28,6 @@ use std::sync::Arc;
 use tracing::*;
 use utils::lsn::Lsn;
 
-use super::bst_layer_map::RetroactiveLayerMap;
-
 ///
 /// LayerMap tracks what layers exist on a timeline.
 ///
@@ -57,11 +55,6 @@ pub struct LayerMap {
     /// All the historic layers are kept here
     historic_layers: RTree<LayerRTreeObject>,
 
-    /// HACK I'm experimenting with a new index to reaplace the RTree. If this
-    ///      works out I'll clean up the struct later.
-    index: RetroactiveLayerMap<Arc<dyn Layer>>,
-    images: RetroactiveLayerMap<Arc<dyn Layer>>,
-
     /// L0 layers have key range Key::MIN..Key::MAX, and locating them using R-Tree search is very inefficient.
     /// So L0 layers are held in l0_delta_layers vector, in addition to the R-tree.
     l0_delta_layers: Vec<Arc<dyn Layer>>,
@@ -248,65 +241,6 @@ impl LayerMap {
     /// layer.
     ///
     pub fn search(&self, key: Key, end_lsn: Lsn) -> Result<Option<SearchResult>> {
-        let old = self.search_old(key, end_lsn)?;
-        let new = self.search_new(key, end_lsn)?;
-        match (&old, &new) {
-            (None, None) => {}
-            (None, Some(_)) => panic!("returned Some, expected None"),
-            (Some(_), None) => panic!("returned None, expected Some"),
-            (Some(old), Some(new)) => {
-                // TODO be more verbose and flexible
-                let context = format!("query: key {}, end_lsn: {}", key, end_lsn);
-                assert_eq!(old.layer.filename(), new.layer.filename(), "{}", context);
-                assert_eq!(old.lsn_floor, new.lsn_floor, "{}", context);
-            }
-        }
-        return Ok(new);
-    }
-
-    // HACK just testing correctness
-    fn search_new(&self, key: Key, end_lsn: Lsn) -> Result<Option<SearchResult>> {
-        // TODO I'm making two separate queries, which is 2x the cost, but that
-        //      can be avoided in varous ways. Caching latest_image queries is
-        //      probably the simplest, but combining the two data structures
-        //      might be better.
-        let latest_layer = self.index.query(key.to_i128(), end_lsn.0 - 1);
-        let latest_image = self.images.query(key.to_i128(), end_lsn.0 - 1);
-
-        // Check for exact match
-        let latest_image = if let Some(image) = latest_image {
-            let img_lsn = image.get_lsn_range().start;
-            if Lsn(img_lsn.0 + 1) == end_lsn {
-                return Ok(Some(SearchResult {
-                    layer: image,
-                    lsn_floor: img_lsn,
-                }));
-            }
-
-            // HACK just to give back ownership of latest_image to parent scope.
-            //      There's definitely a cleaner way to do it.
-            Some(image)
-        } else {
-            None
-        };
-
-        return Ok(latest_layer.map(|layer| {
-            // Compute lsn_floor
-            let mut lsn_floor = layer.get_lsn_range().start;
-            if let Some(image) = latest_image {
-                if layer.is_incremental() {
-                    lsn_floor = std::cmp::max(lsn_floor, image.get_lsn_range().start + 1)
-                }
-            }
-            SearchResult {
-                layer,
-                lsn_floor,
-            }
-        }));
-    }
-
-    // HACK just testing correctness
-    fn search_old(&self, key: Key, end_lsn: Lsn) -> Result<Option<SearchResult>> {
         // linear search
         // Find the latest image layer that covers the given key
         let mut latest_img: Option<Arc<dyn Layer>> = None;
@@ -411,50 +345,19 @@ impl LayerMap {
     /// Insert an on-disk layer
     ///
     pub fn insert_historic(&mut self, layer: Arc<dyn Layer>) {
-        let kr = layer.get_key_range();
-        let lr = layer.get_lsn_range();
-        self.index.insert(
-            kr.start.to_i128()..kr.end.to_i128(),
-            lr.start.0..lr.end.0,
-            Arc::clone(&layer),
-        );
-        if !layer.is_incremental() {
-            self.images.insert(
-                kr.start.to_i128()..kr.end.to_i128(),
-                lr.start.0..lr.end.0,
-                Arc::clone(&layer),
-            );
-        }
-
         if layer.get_key_range() == (Key::MIN..Key::MAX) {
             self.l0_delta_layers.push(layer.clone());
         }
-        // TODO remove this so insert isn't slow. I need it for now for iter_historic()
         self.historic_layers.insert(LayerRTreeObject::new(layer));
         NUM_ONDISK_LAYERS.inc();
     }
 
-    /// Must be called after a batch of insert_historic calls, before querying
-    pub fn rebuild_index(&mut self) {
-        self.index.rebuild();
-        self.images.rebuild();
-    }
-
     ///
     /// Remove an on-disk layer from the map.
     ///
     /// This should be called when the corresponding file on disk has been deleted.
     ///
     pub fn remove_historic(&mut self, layer: Arc<dyn Layer>) {
-        let kr = layer.get_key_range();
-        let lr = layer.get_lsn_range();
-        self.index
-            .remove(kr.start.to_i128()..kr.end.to_i128(), lr.start.0..lr.end.0);
-        if !layer.is_incremental() {
-            self.images
-                .remove(kr.start.to_i128()..kr.end.to_i128(), lr.start.0..lr.end.0);
-        }
-
         if layer.get_key_range() == (Key::MIN..Key::MAX) {
             let len_before = self.l0_delta_layers.len();
 
@@ -683,5 +586,3 @@ impl LayerMap {
         Ok(())
     }
 }
-
-// TODO add layer map tests

pageserver/src/tenant/segment_tree_layer_map.rs

@@ -0,0 +1,346 @@
+use persistent_range_query::naive::{IndexableKey, NaiveVecStorage};
+use persistent_range_query::ops::SameElementsInitializer;
+use persistent_range_query::segment_tree::{MidpointableKey, PersistentSegmentTree, PersistentSegmentTreeVersion};
+use persistent_range_query::{
+    LazyRangeInitializer, PersistentVecStorage, RangeModification, RangeQueryResult,
+    VecReadableVersion,
+};
+use std::cmp::Ordering;
+use std::collections::BTreeMap;
+use std::ops::Range;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
+struct PageIndex(u32);
+type LayerId = String;
+
+impl IndexableKey for PageIndex {
+    fn index(all_keys: &Range<Self>, key: &Self) -> usize {
+        (key.0 as usize) - (all_keys.start.0 as usize)
+    }
+
+    fn element_range(all_keys: &Range<Self>, index: usize) -> Range<Self> {
+        PageIndex(all_keys.start.0 + index as u32)..PageIndex(all_keys.start.0 + index as u32 + 1)
+    }
+}
+
+impl MidpointableKey for PageIndex {
+    fn midpoint(range: &Range<Self>) -> Self {
+        PageIndex(range.start.0 + (range.end.0 - range.start.0) / 2)
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+struct LayerMapInformation {
+    // Only make sense for a range of length 1.
+    last_layer: Option<LayerId>,
+    last_image_layer: Option<LayerId>,
+    // Work for all ranges
+    max_delta_layers: (usize, Range<PageIndex>),
+}
+
+impl LayerMapInformation {
+    fn last_layers(&self) -> (&Option<LayerId>, &Option<LayerId>) {
+        (&self.last_layer, &self.last_image_layer)
+    }
+
+    fn max_delta_layers(&self) -> &(usize, Range<PageIndex>) {
+        &self.max_delta_layers
+    }
+}
+
+fn merge_ranges(left: &Range<PageIndex>, right: &Range<PageIndex>) -> Range<PageIndex> {
+    if left.is_empty() {
+        right.clone()
+    } else if right.is_empty() {
+        left.clone()
+    } else if left.end == right.start {
+        left.start..right.end
+    } else {
+        left.clone()
+    }
+}
+
+impl RangeQueryResult<PageIndex> for LayerMapInformation {
+    fn new_for_empty_range() -> Self {
+        LayerMapInformation {
+            last_layer: None,
+            last_image_layer: None,
+            max_delta_layers: (0, PageIndex(0)..PageIndex(0)),
+        }
+    }
+
+    fn combine(
+        left: &Self,
+        _left_range: &Range<PageIndex>,
+        right: &Self,
+        _right_range: &Range<PageIndex>,
+    ) -> Self {
+        // Note that either range may be empty.
+        LayerMapInformation {
+            last_layer: left
+                .last_layer
+                .as_ref()
+                .or_else(|| right.last_layer.as_ref())
+                .cloned(),
+            last_image_layer: left
+                .last_image_layer
+                .as_ref()
+                .or_else(|| right.last_image_layer.as_ref())
+                .cloned(),
+            max_delta_layers: match left.max_delta_layers.0.cmp(&right.max_delta_layers.0) {
+                Ordering::Less => right.max_delta_layers.clone(),
+                Ordering::Greater => left.max_delta_layers.clone(),
+                Ordering::Equal => (
+                    left.max_delta_layers.0,
+                    merge_ranges(&left.max_delta_layers.1, &right.max_delta_layers.1),
+                ),
+            },
+        }
+    }
+
+    fn add(
+        left: &mut Self,
+        left_range: &Range<PageIndex>,
+        right: &Self,
+        right_range: &Range<PageIndex>,
+    ) {
+        *left = Self::combine(&left, left_range, right, right_range);
+    }
+}
+
+#[derive(Clone, Debug)]
+struct AddDeltaLayers {
+    last_layer: LayerId,
+    count: usize,
+}
+
+#[derive(Clone, Debug)]
+struct LayerMapModification {
+    add_image_layer: Option<LayerId>,
+    add_delta_layers: Option<AddDeltaLayers>,
+}
+
+impl LayerMapModification {
+    fn add_image_layer(layer: impl Into<LayerId>) -> Self {
+        LayerMapModification {
+            add_image_layer: Some(layer.into()),
+            add_delta_layers: None,
+        }
+    }
+
+    fn add_delta_layer(layer: impl Into<LayerId>) -> Self {
+        LayerMapModification {
+            add_image_layer: None,
+            add_delta_layers: Some(AddDeltaLayers {
+                last_layer: layer.into(),
+                count: 1,
+            }),
+        }
+    }
+}
+
+impl RangeModification<PageIndex> for LayerMapModification {
+    type Result = LayerMapInformation;
+
+    fn no_op() -> Self {
+        LayerMapModification {
+            add_image_layer: None,
+            add_delta_layers: None,
+        }
+    }
+
+    fn is_no_op(&self) -> bool {
+        self.add_image_layer.is_none() && self.add_delta_layers.is_none()
+    }
+
+    fn is_reinitialization(&self) -> bool {
+        self.add_image_layer.is_some()
+    }
+
+    fn apply(&self, result: &mut Self::Result, range: &Range<PageIndex>) {
+        if let Some(layer) = &self.add_image_layer {
+            result.last_layer = Some(layer.clone());
+            result.last_image_layer = Some(layer.clone());
+            result.max_delta_layers = (0, range.clone());
+        }
+        if let Some(AddDeltaLayers { last_layer, count }) = &self.add_delta_layers {
+            result.last_layer = Some(last_layer.clone());
+            result.max_delta_layers.0 += count;
+        }
+    }
+
+    fn compose(later: &Self, earlier: &mut Self) {
+        if later.add_image_layer.is_some() {
+            *earlier = later.clone();
+            return;
+        }
+        if let Some(AddDeltaLayers { last_layer, count }) = &later.add_delta_layers {
+            let res = earlier.add_delta_layers.get_or_insert(AddDeltaLayers {
+                last_layer: LayerId::default(),
+                count: 0,
+            });
+            res.last_layer = last_layer.clone();
+            res.count += count;
+        }
+    }
+}
+
+impl LazyRangeInitializer<LayerMapInformation, PageIndex> for SameElementsInitializer<()> {
+    fn get(&self, range: &Range<PageIndex>) -> LayerMapInformation {
+        LayerMapInformation {
+            last_layer: None,
+            last_image_layer: None,
+            max_delta_layers: (0, range.clone()),
+        }
+    }
+}
+
+type Head = PersistentSegmentTree<LayerMapModification, SameElementsInitializer<()>, PageIndex>;
+type Version = PersistentSegmentTreeVersion<LayerMapModification, SameElementsInitializer<()>, PageIndex>;
+
+pub struct STLM {
+    head: Head,
+    historic: BTreeMap<u32, Version>,
+}
+
+/// Layer map (good enough for benchmarks) implemented using persistent segment tree
+impl STLM {
+    pub fn new() -> Self {
+        STLM {
+            head: PersistentSegmentTree::new(
+                PageIndex(0)..PageIndex(100),
+                SameElementsInitializer::new(()),
+            ),
+            historic: BTreeMap::default(),
+        }
+    }
+
+    pub fn insert(self: &mut Self, key_begin: u32, key_end: u32, lsn: u32, value: String) {
+        self.head.modify(
+            &(PageIndex(key_begin)..PageIndex(key_end)),
+            &LayerMapModification::add_image_layer(value),
+        );
+        self.historic.insert(lsn, self.head.freeze());
+    }
+
+    pub fn query(self: &Self, key: u32, lsn: u32) -> Option<String> {
+        let version = self.historic.range(0..=lsn).rev().next()?.1;
+        let info = version.get(&(PageIndex(key)..PageIndex(key + 1)));
+        info.last_image_layer.map(|s| s.clone())
+    }
+}
+
+fn test_stlm() {
+    let mut stlm = STLM::new();
+    stlm.insert(0, 5, 100, "layer 1".to_string());
+    stlm.insert(3, 9, 110, "layer 2".to_string());
+
+    dbg!(stlm.query(1, 105));
+    dbg!(stlm.query(4, 105));
+    dbg!(stlm.query(4, 115));
+}
+
+#[test]
+fn test_stlm_() {
+    test_stlm()
+}
+
+fn test_layer_map<
+    S: PersistentVecStorage<LayerMapModification, SameElementsInitializer<()>, PageIndex>,
+>() {
+    let mut s = S::new(
+        PageIndex(0)..PageIndex(100),
+        SameElementsInitializer::new(()),
+    );
+    s.modify(
+        &(PageIndex(0)..PageIndex(70)),
+        &LayerMapModification::add_image_layer("Img0..70"),
+    );
+    s.modify(
+        &(PageIndex(50)..PageIndex(100)),
+        &LayerMapModification::add_image_layer("Img50..100"),
+    );
+    s.modify(
+        &(PageIndex(10)..PageIndex(60)),
+        &LayerMapModification::add_delta_layer("Delta10..60"),
+    );
+    let s_before_last_delta = s.freeze();
+    s.modify(
+        &(PageIndex(20)..PageIndex(80)),
+        &LayerMapModification::add_delta_layer("Delta20..80"),
+    );
+
+    assert_eq!(
+        s.get(&(PageIndex(5)..PageIndex(6))).last_layers(),
+        (&Some("Img0..70".to_owned()), &Some("Img0..70".to_owned()))
+    );
+    assert_eq!(
+        s.get(&(PageIndex(15)..PageIndex(16))).last_layers(),
+        (
+            &Some("Delta10..60".to_owned()),
+            &Some("Img0..70".to_owned())
+        )
+    );
+    assert_eq!(
+        s.get(&(PageIndex(25)..PageIndex(26))).last_layers(),
+        (
+            &Some("Delta20..80".to_owned()),
+            &Some("Img0..70".to_owned())
+        )
+    );
+    assert_eq!(
+        s.get(&(PageIndex(65)..PageIndex(66))).last_layers(),
+        (
+            &Some("Delta20..80".to_owned()),
+            &Some("Img50..100".to_owned())
+        )
+    );
+    assert_eq!(
+        s.get(&(PageIndex(95)..PageIndex(96))).last_layers(),
+        (
+            &Some("Img50..100".to_owned()),
+            &Some("Img50..100".to_owned())
+        )
+    );
+
+    assert_eq!(
+        s.get(&(PageIndex(0)..PageIndex(100))).max_delta_layers(),
+        &(2, PageIndex(20)..PageIndex(60)),
+    );
+    assert_eq!(
+        *s_before_last_delta
+            .get(&(PageIndex(0)..PageIndex(100)))
+            .max_delta_layers(),
+        (1, PageIndex(10)..PageIndex(60)),
+    );
+
+    assert_eq!(
+        *s.get(&(PageIndex(10)..PageIndex(30))).max_delta_layers(),
+        (2, PageIndex(20)..PageIndex(30))
+    );
+    assert_eq!(
+        *s.get(&(PageIndex(10)..PageIndex(20))).max_delta_layers(),
+        (1, PageIndex(10)..PageIndex(20))
+    );
+
+    assert_eq!(
+        *s.get(&(PageIndex(70)..PageIndex(80))).max_delta_layers(),
+        (1, PageIndex(70)..PageIndex(80))
+    );
+    assert_eq!(
+        *s_before_last_delta
+            .get(&(PageIndex(70)..PageIndex(80)))
+            .max_delta_layers(),
+        (0, PageIndex(70)..PageIndex(80))
+    );
+}
+
+#[test]
+fn test_naive() {
+    test_layer_map::<NaiveVecStorage<_, _, _>>();
+}
+
+#[test]
+fn test_segment_tree() {
+    test_layer_map::<PersistentSegmentTree<_, _, _>>();
+}

pageserver/src/tenant/timeline.rs

@@ -931,7 +931,6 @@ impl Timeline {
                 trace!("found layer {}", layer.filename().display());
                 total_physical_size += layer.path().metadata()?.len();
                 layers.insert_historic(Arc::new(layer));
-                layers.rebuild_index();
                 num_layers += 1;
             } else if let Some(deltafilename) = DeltaFileName::parse_str(&fname) {
                 // Create a DeltaLayer struct for each delta file.
@@ -956,7 +955,6 @@ impl Timeline {
                 trace!("found layer {}", layer.filename().display());
                 total_physical_size += layer.path().metadata()?.len();
                 layers.insert_historic(Arc::new(layer));
-                layers.rebuild_index();
                 num_layers += 1;
             } else if fname == METADATA_FILE_NAME || fname.ends_with(".old") {
                 // ignore these
@@ -1087,11 +1085,10 @@ impl Timeline {
                 let image_layer =
                     ImageLayer::new(self.conf, self.timeline_id, self.tenant_id, &imgfilename);
 
-                {
-                    let mut layers = self.layers.write().unwrap();
-                    layers.insert_historic(Arc::new(image_layer));
-                    layers.rebuild_index();
-                }
+                self.layers
+                    .write()
+                    .unwrap()
+                    .insert_historic(Arc::new(image_layer));
                 self.metrics.current_physical_size_gauge.add(sz);
             } else if let Some(deltafilename) = DeltaFileName::parse_str(fname) {
                 // Create a DeltaLayer struct for each delta file.
@@ -1118,11 +1115,10 @@ impl Timeline {
                 let delta_layer =
                     DeltaLayer::new(self.conf, self.timeline_id, self.tenant_id, &deltafilename);
 
-                {
-                    let mut layers = self.layers.write().unwrap();
-                    layers.insert_historic(Arc::new(delta_layer));
-                    layers.rebuild_index();
-                }
+                self.layers
+                    .write()
+                    .unwrap()
+                    .insert_historic(Arc::new(delta_layer));
                 self.metrics.current_physical_size_gauge.add(sz);
             } else {
                 bail!("unexpected layer filename in remote storage: {}", fname);
@@ -1815,7 +1811,6 @@ impl Timeline {
         {
             let mut layers = self.layers.write().unwrap();
             layers.insert_historic(Arc::new(new_delta));
-            layers.rebuild_index();
         }
 
         // update the timeline's physical size
@@ -1975,7 +1970,6 @@ impl Timeline {
             self.metrics.current_physical_size_gauge.add(metadata.len());
             layers.insert_historic(Arc::new(l));
         }
-        layers.rebuild_index();
         drop(layers);
         timer.stop_and_record();
 
@@ -2282,7 +2276,6 @@ impl Timeline {
 
             new_layer_paths.insert(new_delta_path, LayerFileMetadata::new(metadata.len()));
             layers.insert_historic(Arc::new(l));
-            layers.rebuild_index();
         }
 
         // Now that we have reshuffled the data to set of new delta layers, we can
@@ -2298,7 +2291,6 @@ impl Timeline {
             l.delete()?;
             layers.remove_historic(l);
         }
-        layers.rebuild_index();
         drop(layers);
 
         // Also schedule the deletions in remote storage
@@ -2595,7 +2587,6 @@ impl Timeline {
             layers.remove_historic(doomed_layer);
             result.layers_removed += 1;
         }
-        layers.rebuild_index();
 
         info!(
             "GC completed removing {} layers, cutoff {}",