Add preliminary support for generators with new coroutine system

rt/include/python/builtins.hpp

@@ -48,12 +48,20 @@ concept PyNext = requires(T t) {
   { t.py_next() } -> std::same_as<std::optional<typename T::value_type>>;
 };
 
-template <PyNext T>
-std::optional<typename T::value_type>
-next(T &t, std::optional<typename T::value_type> def = std::nullopt) {
-  auto opt = t.py_next();
-  return opt ? opt : def;
-}
+struct {
+  template <PyNext T>
+  std::optional<typename T::value_type>
+  sync(T &t, std::optional<typename T::value_type> def = std::nullopt) {
+    auto opt = t.py_next();
+    return opt ? opt : def;
+  }
+
+  template <PyNext T>
+  auto operator()(T &t, std::optional<typename T::value_type> def = std::nullopt)
+  -> typon::Task<decltype(sync(t, def))> {
+    co_return sync(t, def);
+  }
+} next;
 
 template <typename T>
 std::ostream &operator<<(std::ostream &os, std::optional<T> const &opt) {

rt/include/python/builtins/range.hpp

@@ -7,11 +7,25 @@
 
 #include <ranges>
 
+#include <typon/typon.hpp>
+
 // todo: proper range support
-template <typename T> auto range(T stop) { return std::views::iota(0, stop); }
+struct {
+  template <typename T> auto sync(T stop) { return std::views::iota(0, stop); }
+
+  template <typename T> auto sync(T start, T stop) {
+    return std::views::iota(start, stop);
+  }
+
+  template <typename T>
+  auto operator()(T stop) -> typon::Task<decltype(sync(stop))> {
+    co_return sync(stop);
+  }
 
-template <typename T> auto range(T start, T stop) {
-  return std::views::iota(start, stop);
-}
+  template <typename T>
+  auto operator()(T start, T stop) -> typon::Task<decltype(sync(start, stop))> {
+    co_return sync(start, stop);
+  }
+} range;
 
 #endif // TYPON_RANGE_HPP

trans/transpiler/__init__.py

 # coding: utf-8
 import ast
 from dataclasses import dataclass, field
-from enum import Enum
+from enum import Enum, auto, Flag
 from itertools import chain, zip_longest
 from typing import *
 
@@ -128,7 +128,7 @@ class NodeVisitor:
                     yield from visitor(node)
                     break
             else:
-                return self.missing_impl(node)
+                yield from self.missing_impl(node)
 
     def missing_impl(self, node):
         raise UnsupportedNodeError(node)
@@ -164,8 +164,8 @@ class SearchVisitor(NodeVisitor):
             elif isinstance(val, ast.AST):
                 yield from self.visit(val)
 
-    def default(self, node):
-        pass
+    def match(self, node) -> bool:
+        return next(self.visit(node), False)
 
 
 class PrecedenceContext:
@@ -180,19 +180,20 @@ class PrecedenceContext:
         self.visitor.precedence.pop()
 
 
-class CoroutineMode(Enum):
-    NONE = 0
-    GENERATOR = 1
-    FAKE = 2
-    TASK = 3
+class CoroutineMode(Flag):
+    SYNC = 1
+    FAKE = 2 | SYNC
+    ASYNC = 4
+    GENERATOR = 8 | ASYNC
+    TASK = 16 | ASYNC
 
 
 # noinspection PyPep8Naming
 @dataclass
 class ExpressionVisitor(NodeVisitor):
     scope: "Scope"
+    generator: CoroutineMode
     precedence: List = field(default_factory=list)
-    generator: CoroutineMode = CoroutineMode.NONE
 
     def visit(self, node):
         if type(node) in SYMBOLS:
@@ -210,13 +211,13 @@ class ExpressionVisitor(NodeVisitor):
         """
         Sets the precedence of the next expression.
         """
-        return ExpressionVisitor(self.scope, [op], generator=self.generator)
+        return ExpressionVisitor(self.scope, self.generator, [op])
 
     def reset(self) -> "ExpressionVisitor":
         """
         Resets the precedence stack.
         """
-        return ExpressionVisitor(self.scope, generator=self.generator)
+        return ExpressionVisitor(self.scope, self.generator)
 
     def visit_Tuple(self, node: ast.Tuple) -> Iterable[str]:
         yield "std::make_tuple("
@@ -262,15 +263,15 @@ class ExpressionVisitor(NodeVisitor):
         if getattr(node, "kwargs", None):
             raise NotImplementedError(node, "kwargs")
         func = node.func
-        if self.generator == CoroutineMode.TASK:
+        # TODO: precedence needed?
+        if CoroutineMode.ASYNC in self.generator:
             yield "co_await "
-        elif self.generator == CoroutineMode.FAKE:
+        elif CoroutineMode.FAKE in self.generator:
             func = ast.Attribute(value=func, attr="sync", ctx=ast.Load())
-        with self.prec_ctx("co_await"):
-            yield from self.prec("()").visit(func)
-            yield "("
-            yield from join(", ", map(self.reset().visit, node.args))
-            yield ")"
+        yield from self.prec("()").visit(func)
+        yield "("
+        yield from join(", ", map(self.reset().visit, node.args))
+        yield ")"
 
     def visit_Lambda(self, node: ast.Lambda) -> Iterable[str]:
         yield "[]"
@@ -346,12 +347,10 @@ class ExpressionVisitor(NodeVisitor):
             yield from self.visit(node.orelse)
 
     def visit_Yield(self, node: ast.Yield) -> Iterable[str]:
-        if self.generator == CoroutineMode.NONE:
-            raise UnsupportedNodeError(node)
-        elif self.generator == CoroutineMode.GENERATOR:
+        if CoroutineMode.GENERATOR in self.generator:
             yield "co_yield"
             yield from self.prec("co_yield").visit(node.value)
-        elif self.generator == CoroutineMode.FAKE:
+        elif CoroutineMode.FAKE in self.generator:
             yield "return"
             yield from self.visit(node.value)
         else:
@@ -451,25 +450,38 @@ class Scope:
 @dataclass
 class BlockVisitor(NodeVisitor):
     scope: Scope
-    generator: CoroutineMode = CoroutineMode.NONE
+    generator: CoroutineMode = CoroutineMode.SYNC
 
     def expr(self) -> ExpressionVisitor:
-        return ExpressionVisitor(self.scope, generator=self.generator)
+        return ExpressionVisitor(self.scope, self.generator)
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> Iterable[str]:
-        yield from self.visit_coroutine(node)
-        # try:
-        #     [*result] = self.visit_func(node, CoroutineMode.NONE)
-        #     return result
-        # except UnsupportedNodeError as e:
-        #     if isinstance(e.node, ast.Yield):
-        #         return self.visit_coroutine(node)
-        #     raise
-
-    def visit_coroutine(self, node: ast.FunctionDef) -> Iterable[str]:
         yield "struct {"
         yield from self.visit_func(node, CoroutineMode.FAKE)
-        yield from self.visit_func(node, CoroutineMode.TASK)
+
+        class YieldVisitor(SearchVisitor):
+            def visit_Yield(self, node: ast.Yield) -> bool:
+                yield True
+
+            def visit_FunctionDef(self, node: ast.FunctionDef):
+                yield from ()
+
+            def visit_ClassDef(self, node: ast.ClassDef):
+                yield from ()
+
+        has_yield = YieldVisitor().match(node.body)
+        yield from self.visit_func(node, CoroutineMode.GENERATOR if has_yield else CoroutineMode.TASK)
+        if has_yield:
+            templ, args, names = self.process_args(node.args)
+            if templ:
+                yield "template"
+                yield templ
+            yield f"auto operator()"
+            yield args
+            yield f"-> typon::Task<decltype(gen({', '.join(names)}))>"
+            yield "{"
+            yield f"co_return gen({', '.join(names)});"
+            yield "}"
         yield f"}} {node.name};"
 
     def visit_func(self, node: ast.FunctionDef, generator: CoroutineMode) -> Iterable[str]:
@@ -482,25 +494,35 @@ class BlockVisitor(NodeVisitor):
             def visit_Return(self, node: ast.Return) -> bool:
                 yield True
 
+            def visit_Yield(self, node: ast.Yield) -> bool:
+                yield True
+
             def visit_FunctionDef(self, node: ast.FunctionDef):
                 yield from ()
 
             def visit_ClassDef(self, node: ast.ClassDef):
                 yield from ()
 
-        has_return = next(ReturnVisitor().visit(node.body), False)
+        has_return = ReturnVisitor().match(node.body)
 
-        if generator == CoroutineMode.FAKE:
+        if CoroutineMode.SYNC in generator:
             if has_return:
                 yield "auto"
             else:
                 yield "void"
             yield "sync"
+        elif CoroutineMode.GENERATOR in generator:
+            yield "auto gen"
         else:
-            yield f"auto operator()"
+            yield "auto operator()"
         yield args
-        if generator == CoroutineMode.TASK:
-            yield f"-> typon::Task<decltype(sync({', '.join(names)}))>"
+        if CoroutineMode.ASYNC in generator:
+            yield "-> typon::"
+            if CoroutineMode.TASK in generator:
+                yield "Task"
+            elif CoroutineMode.GENERATOR in generator:
+                yield "Generator"
+            yield f"<decltype(sync({', '.join(names)}))>"
         yield "{"
         inner_scope = self.scope.function(vars={node.name: VarDecl(VarKind.SELF, None)})
         for child in node.body:
@@ -555,9 +577,9 @@ class BlockVisitor(NodeVisitor):
                 elif decl.kind in (VarKind.GLOBAL, VarKind.NONLOCAL):  # `global` and `nonlocal` just get hoisted as-is.
                     inner_scope.vars[var] = decl
             yield from child_code  # Yeet back the child node code.
-        if generator == CoroutineMode.FAKE:
+        if CoroutineMode.FAKE in generator:
             yield "TYPON_UNREACHABLE();"  # So the compiler doesn't complain about missing return statements.
-        elif generator == CoroutineMode.TASK:
+        elif CoroutineMode.TASK in generator:
             if not has_return:
                 yield "co_return;"
         yield "}"
@@ -682,7 +704,7 @@ class FunctionVisitor(BlockVisitor):
                 yield from self.emit_block(node.orelse)
 
     def visit_Return(self, node: ast.Return) -> Iterable[str]:
-        if self.generator == CoroutineMode.TASK:
+        if CoroutineMode.ASYNC in self.generator:
             yield "co_return "
         else:
             yield "return "