3 files changed, 23 insertions(+), 6 deletions(-)

M type_stack_calc/base/component.py
M type_stack_calc/call/if_call.py
M type_stack_calc/tp_calc.py

@@ 15,6 15,8 @@ from type_stack_calc.ib.plain_component import PlainComponent
 
 from type_stack_calc.util.stack_check import assert_acceptable_vals
 
+class ReturnStop: pass
+
 class BaseComponentSet(BaseNArgs):
     """Sets a variable."""
     with_vars, n = True, 2


@@ 85,6 87,8 @@ class BaseComponentSet(BaseNArgs):
 
         if not getattr(self, 'destructive', None):
             self.destructive = getattr(to, 'destructive', None) and 'setcomp'
+        if component.key == '~r':
+            return [ReturnStop], self
         return [to], self
 
     def __repr__(self): return ".set"

@@ 7,11 7,25 @@
 
 from type_stack_calc.ib.if_block import IfBlock, IfOneBranch
 
-from type_stack_calc.base.component import BaseComponent
+from type_stack_calc.base.component import BaseComponent, ReturnStop
 
 from type_stack_calc.tp.intersection import TpIntersection
 from type_stack_calc.tp.bool import TpBool
 
+def ret_stack_smoosh(a, b):
+    """Combine the return stacks, assuming both are possible."""
+    if len(a) > 0 and a[-1] == ReturnStop:
+        if len(b) > 0 and b[-1] == ReturnStop:  # Both return.
+            return [ReturnStop]
+        else:  # Only `a` returns, `b` goes on as stack.
+            return b
+    elif len(b) > 0 and b[-1] == ReturnStop:  # Only `b` returns.
+        return a
+    else:  # Both go on stack, have to be combined.
+        assert len(a) == len(b), \
+            f"Smooshing stacks failed, differing lengths. ({len(a)} vs {len(b)})\n{a}, {b}"
+        return [ea.smoosh(eb) for ea, eb in zip(a, b)]
+
 class IfCall:
     """Code objects have a `.if` are set to this. Produces an IfBlock if needed."""
     n, fun_tp = 3, 'method'


@@ 51,9 65,8 @@ class IfCall:
             assert c_case is None, "Not definite True/False but not None?"
             *tc_t, if_t = sc.tp_calc(then,      vs, [])
             *tc_f, if_f = sc.tp_calc(otherwise, vs, [])
-            assert len(if_f) == len(if_t), \
-                "Outputs of `().if` must have same stack size."
-            ret_stack = [f.smoosh(t) for f,t in zip(if_t, if_f)]
+                    
+            ret_stack = ret_stack_smoosh(if_t, if_f)
 
             return ret_stack, IfBlock(tc_cond, tc_t, tc_f, ret_stack)
 

@@ 5,7 5,7 @@
 #  by the Free Software Foundation, either version 3 of the License, or
 #  (at your option) any later version.
 
-from type_stack_calc.base.component import BaseComponent
+from type_stack_calc.base.component import BaseComponent, ReturnStop
 from type_stack_calc.ib.scope import ScopeVariable
 
 from type_stack_calc.tp.range import TpRange


@@ 109,7 109,7 @@ class StackCalc:
         code = iter(code)
         while True:
             word = next(code, None)  # Next word, if ran out lastly stack.
-            if word is None:
+            if word is None or len(stack)>0 and stack[-1] == ReturnStop:
                 yield stack
                 return