From 6df9e1ef280353a17f32472afcecccdc56a2eadc Mon Sep 17 00:00:00 2001
From: Jasper den Ouden <o.jasper@gmail.com>
Date: Tue, 8 Mar 2022 01:37:02 +0100
Subject: [PATCH] After returning nothing is to be done, `ReturnStop` does
 this, also logic for this in via `ret_stack_smoosh` in `IfCall`

---
 type_stack_calc/base/component.py |  4 ++++
 type_stack_calc/call/if_call.py   | 21 +++++++++++++++++----
 type_stack_calc/tp_calc.py        |  4 ++--
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/type_stack_calc/base/component.py b/type_stack_calc/base/component.py
index 0d95120..705e08d 100644
--- a/type_stack_calc/base/component.py
+++ b/type_stack_calc/base/component.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"
diff --git a/type_stack_calc/call/if_call.py b/type_stack_calc/call/if_call.py
index 6f73d17..14564ef 100644
--- a/type_stack_calc/call/if_call.py
+++ b/type_stack_calc/call/if_call.py
@@ -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)
 
diff --git a/type_stack_calc/tp_calc.py b/type_stack_calc/tp_calc.py
index 311566b..882da5c 100644
--- a/type_stack_calc/tp_calc.py
+++ b/type_stack_calc/tp_calc.py
@@ -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
 
-- 
2.38.5