Add support for channels

Author: RainingComputers

.vscode/settings.json

@@ -30,6 +30,7 @@
     "Fcomplex",
     "flowargs",
     "fsyntax",
+    "futs",
     "GOARCH",
     "gofmt",
     "goimports",
@@ -105,6 +106,7 @@
     "varsser",
     "venv",
     "vmihailenco",
+    "WAITGROUP",
     "XDECREF"
   ],
   "python.defaultInterpreterPath": "./.venv/bin/python",

pkg/exports.go

@@ -73,6 +73,16 @@ func Send(line int, id func(Context) string, value func(Context) any) AwaitBlock
 func Recv[T any](line int, id func(Context) string, ret func(Context, T)) AwaitBlock {
 	return runtime.Recv[T](line, id, ret)
 }
+func RecvAny(
+	line int,
+	ids []func(Context) string,
+	ref func(Context, []byte, int),
+) AwaitBlock {
+	return runtime.RecvAny(line, ids, ref)
+}
+func Deser[T any](line int, data func(Context) []byte, ret func(Context, T)) SyncBlock {
+	return runtime.Deser(line, data, ret)
+}
 func WaitGroup(
 	line int, value func(Context) int,
 	ret func(Context, string),

pkg/runtimes/dupher/blocks.go

@@ -36,7 +36,7 @@ type BodyLabeledBlock struct {
 }
 type AwaitBlock struct {
 	line int
-	fut  Future
+	futs []Future
 }
 type SyncBlock struct {
 	line int
@@ -84,7 +84,7 @@ func BodyLabeled(line int, label string, stmts ...Block) BodyLabeledBlock {
 }
 
 func Await(line int, start StartFunc, trigger TriggerFunc) AwaitBlock {
-	return AwaitBlock{line: line, fut: Future{start: start, trigger: trigger}}
+	return AwaitBlock{line: line, futs: []Future{{start: start, trigger: trigger}}}
 }
 
 func Sync(line int, stmt func(Context)) SyncBlock {
@@ -294,6 +294,37 @@ func Release(line int, id func(Context) string) AwaitBlock {
 	)
 }
 
+func RecvAny(
+	line int,
+	ids []func(Context) string,
+	ref func(Context, []byte, int),
+) AwaitBlock {
+	futs := make([]Future, len(ids))
+	for idx, id := range ids {
+		futs[idx] = Future{
+			start: func(ctx Context) (platforms.Event, error) {
+				return platforms.ChannelRecvEvent{Id: id(ctx)}, nil
+			},
+			trigger: func(ctx Context, notify platforms.EventNotify) error {
+				recvNotify, ok := notify.(platforms.ChannelRecvEventNotify)
+				if ok {
+					ref(ctx, recvNotify.Value, idx)
+				}
+				return nil
+			},
+		}
+	}
+	return AwaitBlock{line: line, futs: futs}
+}
+
+func Deser[T any](line int, data func(Context) []byte, ret func(Context, T)) SyncBlock {
+	return Sync(line, func(ctx Context) {
+		var v T
+		msgpack.Unmarshal(data(ctx), &v)
+		ret(ctx, v)
+	})
+}
+
 func Program(body Block, source string, name string) ProgramBlock {
 	return ProgramBlock{body: body, Source: source, Name: name}
 }

pkg/runtimes/dupher/compiler.go

@@ -67,9 +67,16 @@ func compileBlock(ctx *compileContext, blk Block) {
 		}
 
 	case AwaitBlock:
+		starts := make([]StartFunc, len(b.futs))
+		triggers := make([]TriggerFunc, len(b.futs))
+		for i, f := range b.futs {
+			starts[i] = f.start
+			triggers[i] = f.trigger
+		}
+
 		ctx.instructions = append(ctx.instructions, Select{
-			Start:   []StartFunc{b.fut.start},
-			Trigger: []TriggerFunc{b.fut.trigger},
+			Start:   starts,
+			Trigger: triggers,
 			Line:    b.line,
 		})
 

pkg/runtimes/dupher/worker_test.go

@@ -544,6 +544,66 @@ func MutexParentFlowSpec() {
 				{sleep: 3 * time.Second, wantVars: map[string]any{"r1": 1, "r2": 1}},
 			},
 		},
+		{
+			name:     "select_channel",
+			flowName: "SelectParentFlow",
+			source: map[string]string{
+				"selectchildfast.ds.go": `package main
+
+import (
+	"github.com/RainingComputers/deepslate/prelude"
+	"sync"
+	"time"
+)
+
+func SelectChildFastSpec(ch chan int, wg sync.WaitGroup) {
+	prelude.Sleep(1 * time.Second)
+	ch <- 1
+	wg.Done()
+}
+`,
+				"selectchildslow.ds.go": `package main
+
+import (
+	"github.com/RainingComputers/deepslate/prelude"
+	"sync"
+	"time"
+)
+
+func SelectChildSlowSpec(ch chan int, wg sync.WaitGroup) {
+	prelude.Sleep(5 * time.Second)
+	ch <- 2
+	wg.Done()
+}
+`,
+				"selectparent.ds.go": `package main
+
+import "sync"
+
+func SelectParentFlowSpec() {
+	ch1 := make(chan int)
+	ch2 := make(chan int)
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go SelectChildFastSpec(ch1, wg)
+	go SelectChildSlowSpec(ch2, wg)
+	var result int = 0
+	select {
+	case x := <-ch1:
+		result = x
+	case y := <-ch2:
+		result = y
+	}
+	_ = result
+	wg.Wait()
+}
+`,
+			},
+			steps: []workerTestStep{
+				{sleep: 3 * time.Second, wantVars: map[string]any{"result": 1}},
+				{sleep: 5 * time.Second},
+			},
+		},
 	}
 
 	httpAddr, controlAddr, shutdown, err := spawnBinary(tests)

pkg/runtimes/dython/deepslate/blocks.py

@@ -67,7 +67,7 @@ def is_block(self) -> None:
 
 @dataclass
 class AwaitBlock:
-    fut: Future
+    futs: list[Future]
     line: int
 
     def is_block(self) -> None:
@@ -139,7 +139,7 @@ def body_labeled(line: int, label: str, *stmts: Block) -> BodyLabeledBlock:
 
 
 def await_block(line: int, start: StartFunc, trigger: TriggerFunc) -> AwaitBlock:
-    return AwaitBlock(fut=Future(start=start, trigger=trigger), line=line)
+    return AwaitBlock(futs=[Future(start=start, trigger=trigger)], line=line)
 
 
 def sync(line: int, stmt: Callable[[Context], None]) -> SyncBlock:
@@ -228,6 +228,39 @@ def _trigger(ctx: Context, notify: EventNotify) -> None:
     return await_block(line, _start, _trigger)
 
 
+def recv_any(
+    line: int,
+    count: int,
+    ids: Callable[[Context], list[str]],
+    ref: Callable[[Context, bytes, int], None],
+) -> AwaitBlock:
+    def _make_future(idx: int) -> Future:
+        def _start(ctx: Context) -> Event:
+            return ChannelRecvEvent(id=ids(ctx)[idx])
+
+        def _trigger(ctx: Context, notify: EventNotify) -> None:
+            if hasattr(notify, "value"):
+                ref(ctx, notify.value, idx)
+
+        return Future(start=_start, trigger=_trigger)
+
+    return AwaitBlock(futs=[_make_future(i) for i in range(count)], line=line)
+
+
+def deser(
+    line: int,
+    data: Callable[[Context], bytes],
+    type_: type,
+    ret: Callable[[Context, Any], None],
+) -> SyncBlock:
+    def _stmt(ctx: Context) -> None:
+        origin = get_origin(type_)
+        value = msgspec.msgpack.decode(data(ctx), type=origin if origin else type_)
+        ret(ctx, value)
+
+    return sync(line, _stmt)
+
+
 def recv(
     line: int,
     id: Callable[[Context], str],

pkg/runtimes/dython/deepslate/compiler.py

@@ -103,10 +103,13 @@ def _compile_block(ctx: _CompileContext, blk: Block) -> None:
             _compile_block(ctx, stmt)
 
     elif isinstance(blk, AwaitBlock):
+        starts = [f.start for f in blk.futs]
+        triggers = [f.trigger for f in blk.futs]
+
         ctx.instructions.append(
             Select(
-                start=[blk.fut.start],
-                trigger=[blk.fut.trigger],
+                start=starts,
+                trigger=triggers,
                 line=blk.line,
             )
         )

pkg/runtimes/dython/deepslate/constants.py

@@ -1,3 +1 @@
-"""Shared constants for deepslate stubs."""
-
 _STUB_MSG = "this is a stub function that is expected to be removed by the transpiler"

pkg/runtimes/dython/deepslate/dpslio.py

@@ -1,5 +1,3 @@
-from typing import Any
-
 from deepslate.constants import _STUB_MSG
 
 
@@ -19,19 +17,19 @@ async def release(self) -> None:
         raise NotImplementedError(_STUB_MSG)
 
 
-class TaskGroup:
-    def create_task(self, coro: object) -> None:
+class WaitGroup:
+    def add(self, n: int) -> None:
         raise NotImplementedError(_STUB_MSG)
 
     async def done(self) -> None:
         raise NotImplementedError(_STUB_MSG)
 
-    async def __aenter__(self) -> "TaskGroup":
-        raise NotImplementedError(_STUB_MSG)
-
-    async def __aexit__(self, *args: Any) -> None:
+    async def wait(self) -> None:
         raise NotImplementedError(_STUB_MSG)
 
 
 def create_task(coro: object) -> None:
     raise NotImplementedError(_STUB_MSG)
+
+
+select = object()

pkg/runtimes/dython/deepslate/imports.py

@@ -355,7 +355,7 @@ def _dython_type(
         if annotation is dpslio.Lock:
             return "ds.MutexRef", ""
 
-        if annotation is dpslio.TaskGroup:
+        if annotation is dpslio.WaitGroup:
             return "ds.WgRef", ""
 
     if inspect.isclass(annotation) and dataclasses.is_dataclass(annotation):