Merge pull request #119 from bksaiki/fixes

Fixes and `tree_sum`

Author: bksaiki

fpy2/interpret/default.py

@@ -840,6 +840,21 @@ def _visit_block(self, block: StmtBlock, ctx: Context):
         for stmt in block.stmts:
             self._visit_statement(stmt, ctx)
 
+    def _cvt_return(self, x: Value):
+        match x:
+            case bool() | Float() | Context():
+                return x
+            case Fraction():
+                return Float.from_rational(x) if is_dyadic(x) else x
+            case tuple():
+                return tuple(self._cvt_return(v) for v in x)
+            case list():
+                for i in range(len(x)):
+                    x[i] = self._cvt_return(x[i])
+                return x
+            case _:
+                raise RuntimeError('unreachable')
+
     def _visit_function(self, func: FuncDef, ctx: Context):
         # process free variables
         for var in func.free_vars:
@@ -851,7 +866,7 @@ def _visit_function(self, func: FuncDef, ctx: Context):
             self._visit_block(func.body, ctx)
             raise RuntimeError('no return statement encountered')
         except FunctionReturnError as e:
-            return e.value
+            return self._cvt_return(e.value)
 
     def _visit_expr(self, e: Expr, ctx: Context) -> Value:
         return super()._visit_expr(e, ctx)

fpy2/libraries/core.py

@@ -234,6 +234,42 @@ def max_e(xs: list[fp.Real]) -> tuple[fp.Real, bool]:
 
     return (largest_e, any_non_zero)
 
+############################################################
+# Arithmetic
+
+@fp.fpy
+def tree_sum(xs: list[fp.Real]):
+    """
+    Sums the elements of xs in a tree.
+    Each sum is rounded under the current rounding context.
+
+    Args:
+        xs: A list of real numbers. The length of xs must be positive and a power of 2.
+
+    Returns:
+        The sum of the elements of xs.
+    """
+
+    with fp.INTEGER:
+        n: fp.Real = len(xs)
+        assert n > 0, "Length of xs must be positive"
+
+        depth = fp.log2(n)
+        assert fp.pow(2, depth) == n, "Length of xs must be a power of 2"
+
+    t = [x for x in xs]
+    for _ in range(depth):
+        with fp.INTEGER:
+            n /= 2
+
+        for i in range(n): # type: ignore[arg-type]
+            with fp.INTEGER:
+                j = 2 * i
+                k = 2 * i + 1
+            t[i] = t[j] + t[k]
+
+    return t[0]
+
 ############################################################
 # Context operations
 

pyproject.toml

@@ -19,6 +19,7 @@ requires-python = ">=3.11"
 dependencies = [
   "titanfp==0.1.1",
   "gmpy2==2.2",
+  "matplotlib"
 ]
 
 [project.optional-dependencies]

tests/infra/backend/cpp.py

@@ -189,6 +189,7 @@ def _test_unit(output_dir: Path, no_cc: bool = False):
     '_modf_spec',
     'isinteger',
     '_ldexp_spec',
+    'tree_sum',
     # eft
     'ideal_2sum',
     'ideal_2mul',