Implement context-modifying align 🧩

2 files changed, 84 insertions, 64 deletions

diff --git a/src/layout/flex.rs b/src/layout/flex.rs
index 877e0bf3..98cca2f9 100644
--- a/src/layout/flex.rs
+++ b/src/layout/flex.rs
@@ -24,7 +24,7 @@ pub struct FlexLayouter {
     stack: StackLayouter,
     usable_width: Size,
     run: FlexRun,
-    cached_glue: Option<Layout>,
+    cached_glue: Option<Size2D>,
 }
 
 /// The context for flex layouting.
@@ -70,7 +70,7 @@ enum FlexUnit {
     /// A unit which acts as glue between two [`FlexUnit::Boxed`] units and
     /// is only present if there was no flow break in between the two
     /// surrounding boxes.
-    Glue(Layout),
+    Glue(Size2D),
 }
 
 struct FlexRun {
@@ -106,8 +106,8 @@ impl FlexLayouter {
         self.units.push(FlexUnit::Boxed(layout));
     }
 
-    /// Add a glue layout which can be replaced by a line break.
-    pub fn add_glue(&mut self, glue: Layout) {
+    /// Add a glue box which can be replaced by a line break.
+    pub fn add_glue(&mut self, glue: Size2D) {
         self.units.push(FlexUnit::Glue(glue));
     }
 
@@ -136,12 +136,7 @@ impl FlexLayouter {
     /// Layout a content box into the current flex run or start a new run if
     /// it does not fit.
     fn layout_box(&mut self, boxed: Layout) -> LayoutResult<()> {
-        let glue_width = self
-            .cached_glue
-            .as_ref()
-            .map(|layout| layout.dimensions.x)
-            .unwrap_or(Size::zero());
-
+        let glue_width = self.cached_glue.unwrap_or(Size2D::zero()).x;
         let new_line_width = self.run.size.x + glue_width + boxed.dimensions.x;
 
         if self.overflows_line(new_line_width) {
@@ -164,32 +159,31 @@ impl FlexLayouter {
             self.flush_glue();
         }
 
-        self.add_to_run(boxed);
+        let dimensions = boxed.dimensions;
+        self.run.content.push((self.run.size.x, boxed));
+
+        self.grow_run(dimensions);
 
         Ok(())
     }
 
-    fn layout_glue(&mut self, glue: Layout) {
+    fn layout_glue(&mut self, glue: Size2D) {
         self.flush_glue();
         self.cached_glue = Some(glue);
     }
 
     fn flush_glue(&mut self) {
         if let Some(glue) = self.cached_glue.take() {
-            let new_line_width = self.run.size.x + glue.dimensions.x;
+            let new_line_width = self.run.size.x + glue.x;
             if !self.overflows_line(new_line_width) {
-                self.add_to_run(glue);
+                self.grow_run(glue);
             }
         }
     }
 
-    fn add_to_run(&mut self, layout: Layout) {
-        let x = self.run.size.x;
-
-        self.run.size.x += layout.dimensions.x;
-        self.run.size.y = crate::size::max(self.run.size.y, layout.dimensions.y);
-
-        self.run.content.push((x, layout));
+    fn grow_run(&mut self, dimensions: Size2D) {
+        self.run.size.x += dimensions.x;
+        self.run.size.y = crate::size::max(self.run.size.y, dimensions.y);
     }
 
     fn finish_run(&mut self) -> LayoutResult<()> {
diff --git a/src/layout/tree.rs b/src/layout/tree.rs
index a1918ded..d125bc84 100644
--- a/src/layout/tree.rs
+++ b/src/layout/tree.rs
@@ -12,6 +12,8 @@ struct TreeLayouter<'a, 'p> {
     stack: StackLayouter,
     flex: FlexLayouter,
     style: Cow<'a, TextStyle>,
+    alignment: Alignment,
+    set_newline: bool,
 }
 
 impl<'a, 'p> TreeLayouter<'a, 'p> {
@@ -27,6 +29,8 @@ impl<'a, 'p> TreeLayouter<'a, 'p> {
                 .. FlexContext::from_layout_ctx(ctx, flex_spacing(&ctx.style))
             }),
             style: Cow::Borrowed(ctx.style),
+            alignment: ctx.alignment,
+            set_newline: false,
         }
     }
 
@@ -34,22 +38,28 @@ impl<'a, 'p> TreeLayouter<'a, 'p> {
     fn layout(&mut self, tree: &SyntaxTree) -> LayoutResult<()> {
         for node in &tree.nodes {
             match node {
-                Node::Text(text) => self.layout_text(text, false)?,
+                Node::Text(text) => {
+                    let layout = self.layout_text(text)?;
+                    self.flex.add(layout);
+                    self.set_newline = true;
+                }
 
                 Node::Space => {
                     // Only add a space if there was any content before.
                     if !self.flex.is_empty() {
-                        self.layout_text(" ", true)?;
+                        let layout = self.layout_text(" ")?;
+                        self.flex.add_glue(layout.dimensions);
                     }
                 }
 
                 // Finish the current flex layouting process.
                 Node::Newline => {
-                    self.layout_flex()?;
+                    self.finish_flex()?;
 
-                    if !self.stack.current_space_is_empty() {
+                    if self.set_newline {
                         let space = paragraph_spacing(&self.style);
                         self.stack.add_space(space)?;
+                        self.set_newline = false;
                     }
 
                     self.start_new_flex();
@@ -69,31 +79,73 @@ impl<'a, 'p> TreeLayouter<'a, 'p> {
 
     /// Finish the layout.
     fn finish(mut self) -> LayoutResult<MultiLayout> {
-        self.layout_flex()?;
+        self.finish_flex()?;
         self.stack.finish()
     }
 
+    /// Layout a function.
+    fn layout_func(&mut self, func: &FuncCall) -> LayoutResult<()> {
+        let mut ctx = self.ctx;
+        ctx.style = &self.style;
+        ctx.shrink_to_fit = true;
+
+        ctx.space.dimensions = self.stack.remaining();
+        ctx.space.padding = SizeBox::zero();
+
+        if let Some(space) = ctx.followup_spaces.as_mut() {
+            *space = space.usable_space();
+        }
+
+        let commands = func.body.layout(ctx)?;
+
+        for command in commands {
+            match command {
+                Command::Layout(tree) => {
+                    self.layout(tree)?;
+                }
+
+                Command::Add(layout) => {
+                    self.finish_flex()?;
+                    self.stack.add(layout)?;
+                    self.set_newline = true;
+                    self.start_new_flex();
+                }
+
+                Command::AddMany(layouts) => {
+                    self.finish_flex()?;
+                    self.stack.add_many(layouts)?;
+                    self.set_newline = true;
+                    self.start_new_flex();
+                }
+
+                Command::SetAlignment(alignment) => {
+                    self.finish_flex()?;
+                    self.alignment = alignment;
+                    self.start_new_flex();
+                }
+
+                Command::SetStyle(style) => {
+                    *self.style.to_mut() = style;
+                }
+            }
+        }
+
+        Ok(())
+    }
+
     /// Add text to the flex layout. If `glue` is true, the text will be a glue
     /// part in the flex layouter. For details, see [`FlexLayouter`].
-    fn layout_text(&mut self, text: &str, glue: bool) -> LayoutResult<()> {
+    fn layout_text(&mut self, text: &str) -> LayoutResult<Layout> {
         let ctx = TextContext {
             loader: &self.ctx.loader,
             style: &self.style,
         };
 
-        let layout = layout_text(text, ctx)?;
-
-        if glue {
-            self.flex.add_glue(layout);
-        } else {
-            self.flex.add(layout);
-        }
-
-        Ok(())
+        layout_text(text, ctx)
     }
 
     /// Finish the current flex layout and add it the stack.
-    fn layout_flex(&mut self) -> LayoutResult<()> {
+    fn finish_flex(&mut self) -> LayoutResult<()> {
         if self.flex.is_empty() {
             return Ok(());
         }
@@ -108,37 +160,11 @@ impl<'a, 'p> TreeLayouter<'a, 'p> {
     fn start_new_flex(&mut self) {
         let mut ctx = self.flex.ctx();
         ctx.space.dimensions = self.stack.remaining();
+        ctx.alignment = self.alignment;
         ctx.flex_spacing = flex_spacing(&self.style);
 
         self.flex = FlexLayouter::new(ctx);
     }
-
-    /// Layout a function.
-    fn layout_func(&mut self, func: &FuncCall) -> LayoutResult<()> {
-        let mut ctx = self.ctx;
-        ctx.style = &self.style;
-        ctx.shrink_to_fit = true;
-
-        ctx.space.dimensions = self.stack.remaining();
-        ctx.space.padding = SizeBox::zero();
-
-        if let Some(space) = ctx.followup_spaces.as_mut() {
-            *space = space.usable_space();
-        }
-
-        let commands = func.body.layout(ctx)?;
-
-        for command in commands {
-            match command {
-                Command::Layout(tree) => self.layout(tree)?,
-                Command::Add(layout) => self.stack.add(layout)?,
-                Command::AddMany(layouts) => self.stack.add_many(layouts)?,
-                Command::ToggleStyleClass(class) => self.style.to_mut().toggle_class(class),
-            }
-        }
-
-        Ok(())
-    }
 }
 
 fn flex_spacing(style: &TextStyle) -> Size {