From 2c08e76e1e9792e0b596dfac558cb7c6b034b245 Mon Sep 17 00:00:00 2001
From: Ben Pfaff <blp@cs.stanford.edu>
Date: Thu, 13 Mar 2025 12:22:08 -0700
Subject: [PATCH] continue to work on pivot table buidling

---
 rust/pspp/src/output/pivot/mod.rs | 85 +++++++++++++++++++++++++++++--
 1 file changed, 81 insertions(+), 4 deletions(-)

diff --git a/rust/pspp/src/output/pivot/mod.rs b/rust/pspp/src/output/pivot/mod.rs
index 7ea269e88d..d7a4cccc64 100644
--- a/rust/pspp/src/output/pivot/mod.rs
+++ b/rust/pspp/src/output/pivot/mod.rs
@@ -413,8 +413,14 @@ impl DimensionBuilder {
         self.hide_all_labels = true;
         self
     }
-    fn build(self, level: usize, top_index: usize, dimension_labels_in_corner: bool) -> Dimension {
+    fn build(
+        mut self,
+        level: usize,
+        top_index: usize,
+        dimension_labels_in_corner: bool,
+    ) -> Dimension {
         let mut leaves = Vec::with_capacity(self.len);
+        self.root.assign_label_depth(dimension_labels_in_corner);
         let root = self
             .root
             .build(dimension_labels_in_corner, None, &mut leaves);
@@ -436,6 +442,9 @@ pub struct GroupBuilder {
     name: Box<Value>,
     children: Vec<CategoryBuilder>,
     show_label: bool,
+    label_depth: usize,
+    extra_depth: usize,
+    show_label_in_corner: bool
 }
 
 impl GroupBuilder {
@@ -444,6 +453,9 @@ impl GroupBuilder {
             name: Box::new(name),
             children: Vec::new(),
             show_label: true,
+            label_depth: 0,
+            extra_depth: 0,
+                show_label_in_corner: false
         }
     }
     pub fn push<T>(&mut self, value: T)
@@ -466,6 +478,39 @@ impl GroupBuilder {
     fn len(&self) -> usize {
         self.children.iter().map(|category| category.len()).sum()
     }
+    fn assign_label_depth(&mut self, dimension_labels_in_corner: bool) {
+        for child in self.children.iter_mut() {
+            child.assign_label_depth(false);
+        }
+        let depth = self
+            .children
+            .iter()
+            .map(|c| c.label_depth())
+            .max()
+            .unwrap_or_default();
+        for child in self.children.iter_mut() {
+            let extra_depth = depth - child.label_depth();
+            if extra_depth > 0 {
+                child.distribute_extra_depth(extra_depth);
+            }
+            child.set_label_depth(depth);
+        }
+        self.show_label_in_corner = self.show_label && dimension_labels_in_corner;
+        self.label_depth = if self.show_label && !self.show_label_in_corner {
+            depth + 1
+        } else {
+            depth
+        };
+    }
+    fn distribute_extra_depth(&mut self, extra_depth: usize) {
+        if self.children.is_empty() {
+            self.extra_depth += extra_depth;
+        } else {
+            for child in self.children.iter_mut() {
+                child.distribute_extra_depth(extra_depth);
+            }
+        }
+    }
     fn build(
         self,
         dimension_labels_in_corner: bool,
@@ -502,6 +547,8 @@ pub enum CategoryBuilder {
     Leaf {
         name: Box<Value>,
         class: Option<Class>,
+        label_depth: usize,
+        extra_depth: usize,
     },
 }
 
@@ -523,12 +570,12 @@ impl CategoryBuilder {
             Self::Group(group) => {
                 Category::Group(group.build(dimension_labels_in_corner, Some(parent), leaves))
             }
-            Self::Leaf { name, class } => {
+            Self::Leaf { name, class, label_depth, extra_depth } => {
                 let leaf = Arc::new(Leaf {
                     parent,
                     name,
-                    label_depth: 0,
-                    extra_depth: 0,
+                    label_depth,
+                    extra_depth,
                     group_index,
                     data_index: leaves.len(),
                     presentation_index: leaves.len(),
@@ -539,6 +586,32 @@ impl CategoryBuilder {
             }
         }
     }
+    fn assign_label_depth(&mut self, dimension_labels_in_corner: bool) {
+        match self {
+            CategoryBuilder::Group(group) => group.assign_label_depth(dimension_labels_in_corner),
+            CategoryBuilder::Leaf { label_depth, .. } => {
+                *label_depth = 1;
+            }
+        }
+    }
+    fn distribute_extra_depth(&mut self, extra: usize) {
+        match self {
+            CategoryBuilder::Group(group) => group.distribute_extra_depth(extra),
+            CategoryBuilder::Leaf { extra_depth, .. } => *extra_depth += extra,
+        }
+    }
+    fn label_depth(&self) -> usize {
+        match self {
+            CategoryBuilder::Group(group) => group.label_depth,
+            CategoryBuilder::Leaf { label_depth, .. } => *label_depth,
+        }
+    }
+    fn set_label_depth(&mut self, depth: usize) {
+        match self {
+            CategoryBuilder::Group(group) => group.label_depth = depth,
+            CategoryBuilder::Leaf { label_depth, .. } => *label_depth = depth,
+        }
+    }
 }
 
 impl From<Value> for CategoryBuilder {
@@ -546,6 +619,8 @@ impl From<Value> for CategoryBuilder {
         Self::Leaf {
             name: Box::new(name),
             class: None,
+            label_depth: 0,
+            extra_depth: 0
         }
     }
 }
@@ -555,6 +630,8 @@ impl From<(Value, Class)> for CategoryBuilder {
         Self::Leaf {
             name: Box::new(name),
             class: Some(class),
+            label_depth: 0,
+            extra_depth: 0
         }
     }
 }
-- 
2.30.2