diff --git a/src/wip/refproof.pysty b/src/wip/refproof.pysty
index a51bf3c..7c30239 100644
--- a/src/wip/refproof.pysty
+++ b/src/wip/refproof.pysty
@@ -75,7 +75,7 @@ __END_OPTIONS_X__
     \end__PACKAGE_PREFIX__saved@subproof
 {
 
-\NewDocumentEnvironment{refproof}{s m o}
+\NewDocumentEnvironment{refproof}{s m !o}
 {
     \stepcounter{__PACKAGE_PREFIX__proofdepth}
     %Save the current claim counter
@@ -124,7 +124,7 @@ __END_OPTIONS_X__
 
 %%% A 'smart' proof environment
 \AtBeginDocument{
-    \RenewDocumentEnvironment{proof}{s o}
+    \RenewDocumentEnvironment{proof}{!s !o}
     {
         \stepcounter{__PACKAGE_PREFIX__proofdepth}
         %Save the current claim counter
@@ -160,11 +160,12 @@ __END_OPTIONS_X__
     }
 }
 
-\NewDocumentEnvironment{subproof}{s o}
+\NewDocumentEnvironment{subproof}{!s !o}
 {
     \stepcounter{__PACKAGE_PREFIX__proofdepth}
     %Save the current claim counter
-    \setcounter{__PACKAGE_PREFIX__save@claim}{\value{claim}}
+    \protected@edef__PACKAGE_MACRO__(dummy@expand){\arabic{claim}}
+    \FSPush{__PACKAGE_PREFIX__save@claim}{__PACKAGE_MACRO__(dummy@expand)}
     \setcounter{claim}{0}
     \def\theclaim{\arabic{claim}}
     __PACKAGE_MACRO__(saved@subproof)[%
@@ -179,6 +180,9 @@ __END_OPTIONS_X__
 {
     __PACKAGE_MACRO__(saved@endproof)
     %Restore previous claim counter
-    \setcounter{claim}{\value{__PACKAGE_PREFIX__save@claim}}
+    \setcounter{claim}{\FSTop{__PACKAGE_PREFIX__save@claim}}
+    \FSPop{__PACKAGE_PREFIX__save@claim}
     \addtocounter{__PACKAGE_PREFIX__proofdepth}{-1}
 }
+
+\newcommand\proofdepth{\arabic{__PACKAGE_PREFIX__proofdepth}}
diff --git a/tests/wip/proof/test.tex b/tests/wip/proof/test.tex
index ac7be22..ac70429 100644
--- a/tests/wip/proof/test.tex
+++ b/tests/wip/proof/test.tex
@@ -35,14 +35,11 @@
 \end{claim}
 
 \begin{proof}
-    \proofdepth
     \begin{claim}
         hi
     \end{claim}
     \begin{proof}
-        \proofdepth
     \end{proof}
-    \proofdepth
     \begin{claim}
         
     \end{claim}