From 3b083d023a8b6c7d4fe6b1140df90c907230b152 Mon Sep 17 00:00:00 2001
From: Titus Wormer <tituswormer@gmail.com>
Date: Mon, 4 Jul 2022 13:45:32 +0200
Subject: Fix opening/closing tag order of nested attention

---
 src/construct/attention.rs |  2 +-
 src/util/edit_map.rs       | 41 +++++++++++++++++++++++++++++------------
 2 files changed, 30 insertions(+), 13 deletions(-)

(limited to 'src')

diff --git a/src/construct/attention.rs b/src/construct/attention.rs
index 95b4079..d4541b4 100644
--- a/src/construct/attention.rs
+++ b/src/construct/attention.rs
@@ -256,7 +256,7 @@ pub fn resolve(tokenizer: &mut Tokenizer) -> Vec<Event> {
                     }
 
                     // Opening.
-                    edit_map.add(
+                    edit_map.add_before(
                         open_event_index,
                         0,
                         vec![
diff --git a/src/util/edit_map.rs b/src/util/edit_map.rs
index 0d58f01..665367a 100644
--- a/src/util/edit_map.rs
+++ b/src/util/edit_map.rs
@@ -66,18 +66,11 @@ impl EditMap {
         }
     }
     /// Create an edit: a remove and/or add at a certain place.
-    pub fn add(&mut self, index: usize, mut remove: usize, mut add: Vec<Event>) {
-        assert!(!self.consumed, "cannot add after consuming");
-
-        if let Some((curr_remove, mut curr_add)) = self.map.remove(&index) {
-            // To do: these might have to be split in several chunks instead
-            // of one, if links in `curr_add` are supported.
-            remove += curr_remove;
-            curr_add.append(&mut add);
-            add = curr_add;
-        }
-
-        self.map.insert(index, (remove, add));
+    pub fn add(&mut self, index: usize, remove: usize, add: Vec<Event>) {
+        add_impl(self, index, remove, add, false);
+    }
+    pub fn add_before(&mut self, index: usize, remove: usize, add: Vec<Event>) {
+        add_impl(self, index, remove, add, true);
     }
     /// Done, change the events.
     pub fn consume(&mut self, events: &mut [Event]) -> Vec<Event> {
@@ -142,3 +135,27 @@ impl EditMap {
         next_events
     }
 }
+
+fn add_impl(
+    edit_map: &mut EditMap,
+    index: usize,
+    mut remove: usize,
+    mut add: Vec<Event>,
+    before: bool,
+) {
+    assert!(!edit_map.consumed, "cannot add after consuming");
+
+    if let Some((curr_remove, mut curr_add)) = edit_map.map.remove(&index) {
+        // To do: these might have to be split in several chunks instead
+        // of one, if links in `curr_add` are supported.
+        remove += curr_remove;
+        if before {
+            add.append(&mut curr_add);
+        } else {
+            curr_add.append(&mut add);
+            add = curr_add;
+        }
+    }
+
+    edit_map.map.insert(index, (remove, add));
+}
-- 
cgit