\ProvidesExplPackage {luamml-patches-amsmath} {2024-10-30} {0.2.0}
  {Feel free to add a description here}

\lua_now:n { require'luamml-amsmath' }

% For all of these changes, the redefinitions appear huge.
% But they are almost identical to the original and only
% add luamml commands in appropriate places, so they would
% mostly disappear if there were enough hooks in amsmath.
\IfPackageAtLeastTF{latex-lab-testphase-math}{2025-01-24}
 {}
 { 
   \PackageInfo{luamml}{patching~\string\start@aligned}
    % aligned and friends
    \cs_set:Npn \start@aligned #1#2 {
      \RIfM@
      \else
        \nonmatherr@ { \begin { \@currenvir } }
      \fi
      \savecolumn@ % Assumption: called inside a group
        \UseTaggingSocket{ math/luamml/annotate/false } {}{ \alignedspace@left }
      \ams@start@box {#1} \bgroup
        \maxfields@ #2 \relax
        \ifnum \maxfields@ > \m@ne
          \multiply \maxfields@ \tw@
          \let \math@cr@@@ \math@cr@@@alignedat
          \alignsep@ \z@skip
        \else
          \let \math@cr@@@ \math@cr@@@aligned
          \alignsep@ \minalignsep
        \fi
        \Let@ \chardef \dspbrk@context \@ne
        \default@tag
        \spread@equation % no-op if already called
        \global \column@ \z@
        \ialign \bgroup
          & \column@plus
            \hfil
            \strut@
            $
              \m@th
              \displaystyle
              {##}
              \UseTaggingSocket{math/luamml/save/nNn}{ {} \displaystyle {mtd}}
            $
            \UseTaggingSocket{math/luamml/mtable/finalizecol}{last}
            \tabskip \z@skip
          & \column@plus
            $
              \m@th
              \displaystyle
              {
                {}
                ##
              }
              \UseTaggingSocket{math/luamml/save/nNn}{ {} \displaystyle {mtd}}
            $
            \UseTaggingSocket{math/luamml/mtable/finalizecol}{last}
            \hfil
            \tabskip\alignsep@
          \crcr
          \ams@return@opt@arg
    }
  \PackageInfo{luamml}{patching~gathered}
     \renewcommand \gathered [1] [c] {
       \RIfM@
       \else
         \nonmatherr@ { \begin {gathered} }
       \fi
       \UseTaggingSocket{ math/luamml/annotate/false } {}{ \alignedspace@left }
       \ams@start@box {#1} \bgroup
         \Let@
         \chardef \dspbrk@context \@ne
         \restore@math@cr
         \spread@equation
         \ialign \bgroup
             \hfil
             \strut@
             $
               \m@th
               \displaystyle
               ##
               \UseTaggingSocket{math/luamml/save/nNn}{ {} \displaystyle {mtd}}
             $
             \UseTaggingSocket{math/luamml/mtable/finalizecol}{last}
             \hfil
           \crcr
             \ams@return@opt@arg
     }
   \PackageInfo{luamml}{patching~\string\endaligned}  
   \cs_set:Npn \endaligned {
         \crcr
         \UseExpandableTaggingSocket{math/luamml/mtable/innertable/save}
       \egroup
       \restorecolumn@
     \egroup
     \UseTaggingSocket{math/luamml/mtable/innertable/finalize}
   }
 }

% gather
\cs_set:Npn \gather@ #1 {
  \ingather@true
  \let \split \insplit@
  \let \tag \tag@in@align
  \let \label \label@in@display
  \chardef \dspbrk@context \z@
  \intertext@ \displ@y@ \Let@
  \let \math@cr@@@ \math@cr@@@gather
  \gmeasure@ {#1}
  \global \shifttag@false
  \tabskip \z@skip
  \global \row@ \@ne
  \halign to \displaywidth \bgroup
      \strut@
      \setboxz@h {
        $
          \m@th
          \displaystyle
          {##}
          \UseTaggingSocket{math/luamml/save/nNn}{ {} \displaystyle {mtd}}
        $
      }
      \UseTaggingSocket{math/luamml/mtable/finalizecol}{box}
      \calc@shift@gather
      \set@gather@field
      \tabskip\@centering
    &
      \setboxz@h {
        \strut@
        {##}
      }
      \dim_compare:nNnTF {0pt} = {
        \box_wd:N \c_zero_int
      }
      { \place@tag@gather }
      {
        \place@tag@gather
        \UseTaggingSocket{math/luamml/mtable/tag/set}
      }
      \tabskip \iftagsleft@
        \gdisplaywidth@
      \else
        \z@skip
      \span \fi
      \crcr
      #1
}

\cs_new_eq:NN \__luamml_amsmath_original_gmeasure:n \gmeasure@
\cs_set:Npn \gmeasure@ #1 {
  \exp_last_unbraced:Nno
    \use_ii_i:nn
    { \luamml_ignore: }
    { \__luamml_amsmath_original_gmeasure:n {#1} }
}
\IfPackageAtLeastTF{latex-lab-testphase-math}{2025-01-24}
 {}
 { \PackageInfo{luamml}{patching~\string\endgather}
   \cs_set:Npn \endgather {
         \math@cr
         \black@ \totwidth@
         \UseExpandableTaggingSocket{math/luamml/mtable/finalize} {gather}
     \egroup
     $$
     \ignorespacesafterend
   }
 }

% align and friends
\cs_set:Npn \align@preamble {
  &
    \hfil
    \strut@
    \setboxz@h {
      \@lign
      $
        \m@th
        \displaystyle
        {##}
        \ifmeasuring@
          \luamml_ignore:
        \else
          \UseTaggingSocket{math/luamml/save/nNn}{ {} \displaystyle {mtd}}
        \fi
      $
    }
    \ifmeasuring@
      \savefieldlength@
    \else
      \UseTaggingSocket{math/luamml/mtable/finalizecol}{box}
    \fi
    \set@field
    \tabskip\z@skip
  &
    \setboxz@h {
      \@lign
      $
      \m@th
      \displaystyle
      {
        {}
        ##
      }
      \ifmeasuring@
        \luamml_ignore:
      \else
        \UseTaggingSocket{math/luamml/save/nNn}{ {} \displaystyle {mtd}}
      \fi
      $
    }
    \ifmeasuring@
      \savefieldlength@
    \else
      \UseTaggingSocket{math/luamml/mtable/finalizecol}{box}
    \fi
    \set@field
    \hfil
    \tabskip\alignsep@
}

\cs_set:Npn \math@cr@@@align {
  \ifst@rred
    \nonumber
  \fi
  \if@eqnsw
    \global \tag@true
  \fi
  \global \advance \row@ \@ne
  \add@amps \maxfields@
  \omit
  \kern -\alignsep@
  \iftag@
    \setboxz@h {
      \@lign
      \strut@
      { \make@display@tag }
    }
    \place@tag
    \UseTaggingSocket{math/luamml/mtable/tag/set}
  \fi
  \ifst@rred
  \else
    \global \@eqnswtrue
  \fi
  \global \lineht@ \z@
  \cr
}

% This was lost anyway, as the latex-lab code overwrites
% the definition again.
\IfPackageAtLeastTF{latex-lab-testphase-math}{2025-01-24}
 {}
 { \PackageInfo{luamml}{patching~\string\maketag@@@}
   \cs_set:Npn \maketag@@@ #1 
    {
      \hbox {
       \m@th
       \normalfont
       #1
       \UseTaggingSocket{math/luamml/mtable/tag/save}
     }
    }
% this handled in latex-lab through \common@align@ending    
    \cs_set:Npn \endalign {
      \math@cr
      \black@ \totwidth@
      \UseTaggingSocket{math/luamml/mtable/finalize} {align}
      \egroup
      \ifingather@
        \restorealignstate@
        \egroup
        \nonumber
        \ifnum0=`{\fi\iffalse}\fi
      \else
        $$
      \fi
      \ignorespacesafterend
    }
}

% For a more interesting one, let's consider multline:
\cs_new_eq:NN \__luamml_amsmath_original_multline:n \multline@
\cs_set:Npn \multline@ #1 {
  \__luamml_amsmath_original_multline:n {
    \ifmeasuring@ \else
      \__luamml_amsmath_set_row_columnalign:n {left}
    \fi
    #1
    \ifmeasuring@ \else
      \__luamml_amsmath_set_row_columnalign:n {right}
    \fi
  }
}

\cs_new_eq:NN \__luamml_amsmath_original_mmeasure:n \mmeasure@
\cs_set:Npn \mmeasure@ #1 {
  \exp_last_unbraced:Nno
    \use_ii_i:nn
    { \luamml_ignore: }
    { \__luamml_amsmath_original_mmeasure:n {#1} }
}

% Luckily, {multline} uses \endmultline@math in exactly
% the spot where we have to set the flag.
% Less luckily, \endmultline@math sometimes get overwritten for the last line.
% But that isn't a problem since we want special behavior there anyway.
\cs_set:Npn \endmultline@math {
  \luamml_save:nNn {} \displaystyle {mtd}
  $
  \UseTaggingSocket{math/luamml/mtable/finalizecol}{last}
}

\cs_set:Npn \rendmultline@ {
    \iftag@
      \luamml_save:nNn {} \displaystyle {mtd}
      $
      \UseTaggingSocket{math/luamml/mtable/finalizecol}{last}
      \let \endmultline@math \relax
      \ifshifttag@
        \hskip \multlinegap
        \llap {
          \vtop {
            \raise@tag
            \normalbaselines
            \setbox \@ne \null
            \dp \@ne \lineht@
            \box \@ne
            \hbox {
              \strut@
              \make@display@tag
            }
          }
        }
      \else
        \hskip \multlinetaggap
        \make@display@tag
      \fi
      \UseTaggingSocket{math/luamml/mtable/tag/set}
    \else
      \hskip \multlinegap
    \fi
    \hfilneg
    \math@cr
    \UseExpandableTaggingSocket {math/luamml/mtable/finalize} {multline}
  \egroup
  $$
}

\cs_set:Npn \lendmultline@ {
    \hfilneg
    \hskip\multlinegap
    \math@cr
    \UseExpandableTaggingSocket {math/luamml/mtable/finalize} {multline}
    %\__luamml_amsmath_finalize_table:n {multline}
  \egroup
  $$
}

% Finally some slightly different stuff.
% While {matrix} is covered by {array}, we still have {smallmatrix}:
\renewenvironment {smallmatrix} {
  \UseTaggingSocket{ math/luamml/annotate/false } {} { \null\, }
  \vcenter \bgroup
    \Let@
    \restore@math@cr
    \default@tag
    \baselineskip 6 \ex@
    \lineskip 1.5 \ex@
    \lineskiplimit \lineskip
    \ialign \bgroup
        \hfil
        $
        \m@th
        \scriptstyle
        ##
        % No \scriptsize here since we want to add the mstyle nodes
        \UseTaggingSocket{math/luamml/save/nn}{ {} {mtd}}
        $
        \UseTaggingSocket{math/luamml/mtable/finalizecol}{last}
        \hfil
      &&
        \thickspace
        \hfil
        $
        \m@th
        \scriptstyle
        ##
        % No \scriptsize here since we want to add the mstyle nodes
        \UseTaggingSocket{math/luamml/save/nn}{ {} {mtd}}        
        $
        \UseTaggingSocket{math/luamml/mtable/finalizecol}{last}
        \hfil
      \crcr
}{%
      \crcr
      \UseExpandableTaggingSocket{math/luamml/mtable/smallmatrix/save}
    \egroup
  \egroup
  \UseTaggingSocket{math/luamml/mtable/innertable/finalize}
  \UseTaggingSocket{ math/luamml/annotate/false } {}{ \, }
}

% {cases} is defined by the kernel, but we patch the overwritten version by amsmath.
\cs_set:Npn \env@cases {
  \let \@ifnextchar \new@ifnextchar
  \left \lbrace
    \def \arraystretch {1.2}
    \array {@{}l@{\quad \luamml_ignore:}l@{}}
}


\cs_set:Npn \bBigg@ #1 #2 {
  {
    \ensuremath {
      \Uvextensible height~#1 \big@size axis~exact~#2
    }
  }
}
