amsmath - redefine align to support \abovedisplayshortskip and \belowdisplayshortskip?

OK, I made a thing. It seems to work with align & co. I have no idea what other things it might break. Any thoughts?

\documentclass{article}
\usepackage[tbtags]{mathtools}
\usepackage{regexpatch}
\makeatletter
% a new glue to store the amount we need to modify at the end of the display
\newskip\belowdisplayskip@mod
% a new box to store the left-hand side of the first line inside a split environment
\newbox\split@firstline@LHS
% redefine \start@align to incorporate a call to the check whether the preceding line is short enough
\xpatchcmd{\start@align}{%
    \collect@body\next@
}{%
    \collect@body\next@
    \shortskip@check
}{}{\PackageWarningNoLine{methtools}{Failed to patch \protect\start@align}}
% redefine \endalign to correct the belowskip
\xpatchcmd{\endalign}{%
    $$%
}{%
    $$%
    \vskip\belowdisplayskip@mod
    \global\belowdisplayskip@mod\z@
    \noindent
}{}{\PackageWarningNoLine{methtools}{Failed to patch \protect\endalign}}
% Redefine \intertext@ to correct the belowskip. How we can patch this depends on whether mathtools has been loaded, hence we defer this action until \begin{document}.
\AtBeginDocument{%
    \@ifpackageloaded{mathtools}{\MHInternalSyntaxOn}{\makeatletter}%
    \xpatchcmd{\intertext@}{%
        \penalty\postdisplaypenalty\vskip\belowdisplayskip
    }{%
        \penalty\postdisplaypenalty\vskip\belowdisplayskip
        \vskip\belowdisplayskip@mod
        \global\belowdisplayskip@mod\z@
    }{}{\PackageWarningNoLine{methtools}{Failed to patch \protect\intertext@}}
    \@ifpackageloaded{mathtools}{\MHInternalSyntaxOff}{\makeatother}%
}
% There is some weird behavior if split is used and the "tbtags" option is active. The first line of the split environment seems to be a dummy line set without any width on the left-hand side. Since shortskips need to be determined depending on the width of the first line, many split environments would erroneously be detected as qualifying for shortskips. To work around this we modify \insplit@ to explicitly measure the width of the left-hand side of the first line and further modify \rendsplit@ to insert an empty box of said width.
\xpatchcmd{\insplit@}{% Patching a command containing ## seems to require regexpatch's xpatchcmd.
    \ialign\bgroup
        \hfil
        \strut@
        $\m@th\displaystyle{##}$%
}{%
    \setbox\split@firstline@LHS\box\voidb@x
    \ialign\bgroup
        \hfil
        \strut@
        \setbox\@tempboxa\hbox{$\m@th\displaystyle{##}$}\unhcopy\@tempboxa
        \ifvoid\split@firstline@LHS
            \global\setbox\split@firstline@LHS\box\@tempboxa
        \fi
}{}{\PackageWarningNoLine{methtools}{Failed to patch \protect\insplit@}}
\xpatchcmd{\rendsplit@}{%
    \global\@tempcnta\column@
}{%
    \global\@tempcnta\column@
    \hbox to\wd\split@firstline@LHS{}%
}{}{\PackageWarningNoLine{methtools}{Failed to patch \protect\rendsplit@}}
% check whether shortskips should be used; if yes, immediately account for the aboveskip and save belowskip for later
\def\shortskip@check{%
    \ifmeasuring@\else
        \noalign{\shortskip@calc@shift@lc}%
        \ifdim\dimexpr\eqnshift@+\@tempdimb\relax > \predisplaysize
            \noalign{\vskip\glueexpr\abovedisplayshortskip-\abovedisplayskip\relax}%
            \noalign{\global\belowdisplayskip@mod=\glueexpr\belowdisplayshortskip-\belowdisplayskip\relax}%
        \fi
    \fi
}
% first helper function to calculate how much the first equation row is indented. amsmath documentation:
%    \cs{@tempdimb} will be the ``indentation'' of leftmost end of
%    text, ignoring the \cs{tabskip} glue, i.e., it will be the sum of
%    the maximum widths of any fields to the left of the first
%    non-empty field, plus whatever empty space there is at the
%    beginning of the first non-empty field.
\def\shortskip@calc@shift@lc{% based on amsmath \x@calc@shift@lc
    \row@\@ne
    \column@\z@
    \@tempdima\z@
    \@tempdimb\z@
    \edef\@tempb{\fieldlengths@\row@}%
    \@for\@tempa:=\@tempb\do{%
        \advance\column@\@ne
        \shortskip@lcalc@width
    }%
}
% second helper function to calculate how much the first equation row is indented. amsmath documentation:
%    This macro calculates the ``indentation'' of the current row, as
%    defined above under the description of \cs{x@calc@shift@lc}.
%    This macro is called for each field of the current line, with
%    \cs{@tempa} set to the width of the current field.  Ideally, the
%    loop enclosing \cs{x@lcalc@width} would terminate as soon as
%    \cs{@tempa} is non-zero, but that would be a bit tricky to
%    arrange.  Instead, we use \cs{@tempdima} as a flag to signal when
%    we've encountered the first non-empty field.
\def\shortskip@lcalc@width{% based on amsmath \x@lcalc@width
    \ifdim\@tempdima = \z@
        \ifdim\@tempa > \z@
            \@tempdima\p@
            \ifodd\column@
                \advance\@tempdimb \maxcol@width\column@
                \global\advance\@tempdimb-\@tempa
            \fi
            \count@\column@
            \advance\count@\m@ne
            \divide\count@\tw@
            \advance\@tempcnta-\count@
            \advance\@tempcntb-\count@
        \else
            \global\advance\@tempdimb \maxcol@width\column@\relax
        \fi
    \fi
}
\makeatother
\begin{document}%
If there is just a short line of text before an \verb|equation| environment, then the short skips (\verb|\abovedisplayshortskip| and \verb|\belowdisplayshortskip|) will be used%
\begin{equation}%
  a=b
\end{equation}%
My modifed \verb|align| will now also insert short skips if the preceding line is short enough, such as in this example%
\begin{align}%
  a&=b\\
  a&=b
\end{align}%
It should depend only on the first equation; here shortskips are used:\\short%
\begin{align}%
                        a&=b\\
  aaaaaaaaaaaaaaaaaaaaaaa&=bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
\end{align}%
but not here, since the first equation is too close to the text (less than 2em):\\regular%
\begin{align}%
  a&=b\\
  a&=bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
\end{align}%
It should also work with multiple columns\\short%
\begin{align}%
        a&=b      &      a&=b       &      a&=b\\
  aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb
\end{align}%
It should also work with multiple columns, now the other way round\\regular%
\begin{align}%
  aaaaaa&=bbbbbbb &       a&=b      &      a&=b\\
       a&=b       & aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb
\end{align}%
Let's see what happens if the first entry in the first column is empty\\short%
\begin{align}%
        &         & aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb\\
  aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb
\end{align}%
%
\newpage\noindent
%
It has also redefined the friends such has \verb|align*|\\short%
\begin{align*}%
       a&=b       & aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb\\
  aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb
\end{align*}%
It has also redefined the friends such has \verb|align*|\\regular%
\begin{align*}%
  aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb\\
       a&=b       & aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb
\end{align*}%
What about \verb|alignat|? Seems to work as well\\short%
\begin{alignat}{3}%
  aaaaaa&=bbbbbbb &\quad aaaaaa&=bbbbbbb &\quad aaaaaa&=bbbbbbb\\
  aaaaaa&=bbbbbbb &      aaaaaa&=bbbbbbb &      aaaaaa&=bbbbbbb
\end{alignat}%
What about \verb|alignat|? Seems to work as well\\regular%
\begin{alignat}{3}%
  aaaaaa&=bbbbbbb &       aaaaaa&=bbbbbbb &       aaaaaa&=bbbbbbb\\
  aaaaaa&=bbbbbbb &\qquad aaaaaa&=bbbbbbb &\qquad aaaaaa&=bbbbbbb
\end{alignat}%
There's also \verb|flalign|. Even that one can become short if first entry is missing\\short%
\begin{flalign}%
        &         & aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb\\
  aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb & aaaaaa&=bbbbbbb
\end{flalign}%
It works. Yay. What do you think?\\[\baselineskip]
\begin{tabular}{@{}ll}%
  \verb|\abovedisplayshortskip| & \the\abovedisplayshortskip\\
  \verb|\abovedisplayskip|      & \the\abovedisplayskip\\
  \verb|\belowdisplayshortskip| & \the\belowdisplayshortskip\\
  \verb|\belowdisplayskip|      & \the\belowdisplayskip
\end{tabular}%
%
\newpage\noindent
%
Now let's use \verb|split| inside \verb|align|, where the preceding line is\\short%
\begin{align}%
\begin{split}%
  aaaaaaaaaaa&=bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb\\
             &=ddddddddddddddddddddddddddddddddddd
\end{split}%
\end{align}%
Now let's use \verb|split| inside \verb|align|, where the preceding line is\\not so short%
\begin{align}%
\begin{split}%
  aaaaaaaaaaa&=bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb\\
             &=ddddddddddddddddddddddddddddddddddd
\end{split}%
\end{align}%
Now let's use \verb|split| inside \verb|align|, but switch it around\\still short enough%
\begin{align}%
\begin{split}%
             &=ddddddddddddddddddddddddddddddddddd\\
  aaaaaaaaaaa&=bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
\end{split}%
\end{align}%
It works. Yay. What do you think?
\end{document}%