refactor to gain access to surrounding elements

1 files changed, 33 insertions, 26 deletions

diff --git a/_extensions/js/multicol-nohead/multicol-nohead.lua b/_extensions/js/multicol-nohead/multicol-nohead.lua
index 2b98798..a5bba1b 100644
--- a/_extensions/js/multicol-nohead/multicol-nohead.lua
+++ b/_extensions/js/multicol-nohead/multicol-nohead.lua
@@ -7,41 +7,48 @@ local multicol_end = pandoc.RawBlock('latex', [[
 \end{multicols}
 ]])
 
-local multicol_now = false
-
 -- TODO: include header level 2 and author paragraph above multicolumn
 function is_author(elem)
     return elem.t == Para and elem.content[1].text == "Forfatter:"
 end
 
-function multicol_delimit(elem)
-    if multicol_now == true then
-      return {multicol_end, elem, multicol_begin}
-    else
-      multicol_now = true
-      return {elem, multicol_begin}
-    end
-end
-
-function Header(elem)
+function Pandoc(doc)
   if FORMAT:match 'latex' then
-    if elem.level == 1 then
-      return multicol_delimit(elem)
+
+    -- locate topmost-only elements where multicol state should change
+    local multicol_now = false
+    local multicol_places = {}
+    for i, elem in ipairs(doc.blocks) do
+      if elem.t == "Header" and elem.level == 1 then
+        multicol_places[i] = multicol_now and "renew" or "begin"
+        multicol_now = true
+      elseif elem.t == "HorizontalRule" then
+        if multicol_now then
+          multicol_places[i] = "end"
+          multicol_now = false
+        end
+      end
+    end
+    if multicol_now then
+      multicol_places[#doc.blocks] = "end"
     end
-  end
-end
 
-function HorizontalRule(elem)
-  if FORMAT:match 'latex' then
-    multicol_now = false
-    return {multicol_end, elem}
-  end
-end
+    -- apply multicol, in reverse order since it shifts later indices
+    for i = #doc.blocks, 1, -1 do
+      local state = multicol_places[i]
+      if state ~= nil then
+        if state == "begin" then
+          doc.blocks:insert(i + 1, multicol_begin)
+        elseif state == "renew" then
+          doc.blocks:insert(i + 1, multicol_begin)
+          doc.blocks:insert(i, multicol_end)
+        elseif state == "end" then
+          doc.blocks:insert(i, multicol_end)
+        end
+      end
+    end
 
-function Pandoc(doc)
-  if FORMAT:match 'latex' then
-    if multicol_now == true then
-      doc.blocks:insert(multicol_end)
+    if #multicol_places >= 1 then
       return doc
     end
   end