diff --git a/luamml-amsmath.lua b/luamml-amsmath.lua
index 453b543..eacb91f 100644
--- a/luamml-amsmath.lua
+++ b/luamml-amsmath.lua
@@ -33,7 +33,7 @@ lua.get_functions_table()[funcid] = function()
     spacing[#spacing+1] = n.width == 0 and '0' or '.8em'
   end
   mml_table.columnspacing = table.concat(spacing, ' ', 2, #spacing-2)
-  save_result(mml_table, 0)
+  save_result(mml_table, true)
 end
 
 local last_tag
diff --git a/luamml-tex.lua b/luamml-tex.lua
index c76d1f7..36edff8 100644
--- a/luamml-tex.lua
+++ b/luamml-tex.lua
@@ -36,12 +36,13 @@ end
 --    Bit 2: Integrate with table mechanism
 
 local mlist_buffer
-local mlist_result
+local mlist_result, mlist_display
 
-local function save_result(xml, style)
-  mlist_result = write_xml(make_root(xml, style))
+local function save_result(xml, display)
+  mlist_result, mlist_display = xml, display
   if tex.count.tracingmathml > 1 then
-    texio.write_nl(mlist_result .. '\n')
+    -- Here xml gets wrapped in an mrow to avoid modifying it.
+    texio.write_nl(write_xml(make_root({[0] = 'mrow', xml}, display and 0 or 2)) .. '\n')
   end
 end
 
@@ -64,9 +65,9 @@ luatexbase.add_to_callback('pre_mlist_to_hlist_filter', function(mlist, style)
   end
   local xml = process_mlist(new_mlist, style == 'display' and 0 or 2)
   if flag & 2 == 0 then
-    save_result(xml, (style == 'display' or flag & 1 == 1) and 0 or 2)
-  else
-    assert(style == 'text')
+    save_result(xml, style == 'display' or flag & 1 == 1)
+  end
+  if style == 'text' then
     local startmath = tex.nest.top.tail
     local props = properties[startmath]
     if not props then
@@ -90,10 +91,11 @@ lua.get_functions_table()[funcid] = function()
         "I was asked to provide MathML code for the last formula, but there weren't any new formulas since you last asked."
       })
   end
+  local mml = write_xml(make_root(mlist_result, mlist_display and 0 or 2))
   if tex.count.tracingmathml == 1 then
-    texio.write_nl(mlist_result .. '\n')
+    texio.write_nl(mml .. '\n')
   end
-  tex.sprint(-2, tostring(pdf.immediateobj('stream', mlist_result, '/Subtype/application#2Fmathml+xml' .. token.scan_argument(true))))
+  tex.sprint(-2, tostring(pdf.immediateobj('stream', mml, '/Subtype/application#2Fmathml+xml' .. token.scan_argument(true))))
   mlist_result = nil
 end