Locally adjust a dimension written into aux

Question

I am writing some code that makes use of a dimension (termed \shape below) that gets adjusted several times over the course of the code. The value \shape has by the time the code is finished is what is used when \shape is called. (See my previous question Make code retrieve a dimensional value from later on in the code? )

Is there a way to adjust the following code somewhat so that different parts of it are cordoned off from the rest, as it were, and within these parts, \shape gets adjusted locally, as opposed to globally, and the value \shape has when called within such a part is the value it has by the end of that part of the code (as opposed to by the end of the (entire) code)?

\documentclass{article}
\usepackage{tabularx}
\usepackage{showframe}
\setlength{\parindent}{0pt}
% The following ~dozen lines are taken from egreg's answer (https://tex.stackexchange.com/a/694439/277990) to a previous question of mine:
\newdimen\shape
\global\shape=0pt
\makeatletter
\AtEndDocument{\write@auxout{\global\shape=\the\shape}}
\makeatother
\newcommand{\shifter}[1]{%
  \settowidth{\dimen0}{#1}%
  \ifdim\shape<\dimen0
    \global\shape=\dimen0
  \fi
}
\newcommand{\proofpiece}[1]{%
\begin{tabularx}{\linewidth}{r >$r<$ @{} >{\raggedright${}}X<{$} p{\shape}}
#1
\end{tabularx}
}%
\newcommand{\rcc}[1]{% ''rightmost column contents''
#1 \shifter{#1}
}
\begin{document}
The following has the output I desire, and the source code is satisfactory:
\proofpiece{

& \multicolumn{2}{l}{Math}
             & \rcc{Words}\
& m & = 9/3  & \rcc{Words}\
&   & = 3    & \rcc{Words}\

}
\proofpiece{
4. & 34x+2 & = 30/2 & \rcc{Wider words}\
5. &       & = 15   & \rcc{Words}\
}
\vspace{2cm}
However, suppose we write a second proof in the same document:
\proofpiece{

& \multicolumn{2}{l}{Math}
             & \rcc{Words}\
& m & = 2+2  & \rcc{Words}\
&   & = 4    & \rcc{Words}\

}
\proofpiece{
4. & 34x+2 & = 16/2 & \rcc{Words}\
5. &       & = 8    & \rcc{Words}\
}
Unfortunately, the width of the rightmost column is equal to the width of the widest entry in the {\em previous} proof (namely, ``Wider words'').
\end{document}

Note: I doubt that this is relevant, but David Carlisle's answer (https://tex.stackexchange.com/a/693585/277990) to one of my previous questions provides context as to why \proofpiece is the way it is.

Answer 1

Use a different dimen register for each group, also beware underfull box warnings from missing %

\documentclass{article}
\usepackage{tabularx}
\usepackage{showframe}
\setlength{\parindent}{0pt}
% The following ~dozen lines are taken from egreg's answer (https://tex.stackexchange.com/a/694439/277990) to a previous question of mine:
\makeatletter
\newdimen\shapea
\AtEndDocument{\immediate\write@auxout{\global\shapea=\the\shapea}}
\newdimen\shapeb
\AtEndDocument{\immediate\write@auxout{\global\shapeb=\the\shapeb}}
\makeatother
\newcommand{\shifter}[1]{%
  \settowidth{\dimen0}{#1}%
  \ifdim\shape<\dimen0
    \global\shape=\dimen0
  \fi
}
\newcommand{\proofpiece}[2]{%
\let\shape#1%
\begin{tabularx}{\linewidth}{r >$r<$ @{} >{\raggedright${}}X<{$} p{#1}}
#2
\end{tabularx}%
}%
\newcommand{\rcc}[1]{% ''rightmost column contents''
#1 \shifter{#1}
}
\begin{document}
The following has the output I desire, and the source code is satisfactory:
\proofpiece\shapea{

& \multicolumn{2}{l}{Math}
             & \rcc{Words}\
& m & = 9/3  & \rcc{Words}\
&   & = 3    & \rcc{Words}\

}
\proofpiece\shapea{
4. & 34x+2 & = 30/2 & \rcc{Wider words}\
5. &       & = 15   & \rcc{Words}\
}
\vspace{2cm}
However, suppose we write a second proof in the same document:
\proofpiece\shapeb{

& \multicolumn{2}{l}{Math}
             & \rcc{Words}\
& m & = 2+2  & \rcc{Words}\
&   & = 4    & \rcc{Words}\

}
\proofpiece\shapeb{
4. & 34x+2 & = 16/2 & \rcc{Words}\
5. &       & = 8    & \rcc{Words}\
}
Unfortunately, the width of the rightmost column is equal to the width of the widest entry in the {\em previous} proof (namely, ``Wider words'').
\end{document}