Hi Rahguzar,

I thought about the further edge cases you raised.  The attached patch
is what I came up with.

You mentioned the following problem: when folding \begin macros, you
want to consume optional arguments following (e.g.) \begin{theorem} but
not (e.g.) \begin{equation}, so as to avoid errors in, e.g.,

\begin{equation}
  [X, Y] = Z
\end{equation}

With the attached patch, you can do so as follows:

--8<---------------cut here---------------start------------->8---
(defun my-equation-env-p (args)
  "Return non-nil when ARGS describes an equation environment.
Return non-nil if the list ARGS is of the form

  '(\"{env}\")

where ENV is an equation environment, such as equation, align, gather,
multline or their starred variants."
  (and (= 1 (length args))
       (let* ((envs (LaTeX--math-environment-list))
              (re (concat "{" (regexp-opt envs) "}")))
         (string-match-p re (car args)))))

(setq-local TeX-fold-macro-spec-list
            `((("⬖ {1}" . my-equation-env-p) ("begin")) 
              (("⬗ {1}" . 1) ("end"))))
--8<---------------cut here---------------end--------------->8---

This says that when determining the folding extent of a \begin macro for
a display math environment, we stop after the first argument (which is
just the environment name).

As before, the "1" signature in the \end macro has the effect that in

\end{equation}
[blah]

we do not fold the trailing "[blah]".

Edge cases such as

\begin{equation}
  \label{eq:blah} [X, Y] = Z
\end{equation}

are a bit subtler: with some packages, \label can accept optional
arguments, so we can't avoid such cases simply by restricting the number
of optional and required arguments.  Instead, we use the macro spec list
entry

(("[l]" . TeX-fold-stop-after-first-required) ("label"))

with the following predicate:

(defun TeX-fold-stop-after-first-required (args)
  "Return nil when final element of ARGS starts with \"{\"."
  (and args (string-prefix-p "{" (car (last args)))))

With this, when we fold \label macros, we stop consuming after the first
required argument that we encounter.  Thus, \label[...][...]{...}  folds
correctly, while \label[...]{...}[...] stops after the {...}.

I can't think of any other edge cases.  I think this patch is now ready
to ship, although I'd welcome any feedback.

Thanks, best,

Paul