Carry waits keyed by a program's own type beside the generated loop

CHANGELOG.md

@@ -54,6 +54,7 @@ The generated loop is now written from the program's source alone, and the manif
 
 ### Fixed
 
+- **A program with the generated loop can key its own waits by a type of its own.** The loop keys its wait by `Int`, and a hand-written `Wait.poll` keyed by a sum beside it used to be refused by the Rust door ("this program keys one wait set by 'Int' and another by 'Watch'") and to fail wasm-gc validation. In a program that answers a capability of its own, each such wait, in the entry or in a dependency, now goes through helpers generated for its key type that carry it through an `Int`-keyed wait. It answers the same keys in the same order. Its recording holds the `Int`-keyed wait. The wasm-gc wait ABI is unchanged, and `--target wasip2` runs such waits too. A program that answers no capability still keys all its waits one way.
 - **`check` no longer asks a verify block of a function no verify case can call.** A parameter of a capability resource type (`Tcp.Connection`, `Work.Job`, a job kind's handle), or of a tuple, record or sum of the module that always carries one, has no value a case can write, so such a pure branching helper failed `error[missing-verify]` with no way to satisfy it. It is now exempt, the way effectful functions are. A parameter with an empty value (`List`, `Option`, a sum with a resource-free variant) still needs its verify block.
 - **A `main` that answers `Err` exits non-zero on wasm-gc and wasip2**, with the error on stderr on wasm-gc, as it already did on the VM and in generated Rust. Both wasm targets used to exit zero.
 - **`aver replay` of a run whose `main` answered `Err` matches.** The VM replay compared a runtime error against the recorded `Err` value and always reported a mismatch.

docs/services.md

@@ -488,7 +488,7 @@ fn interests(eye: Eye, job: Work.Job) -> Map<Watch, Wait.Item>
 
 The key type of a wait is the key type of the map it was handed, and the type checker settles it. A wait set written at the call is keyed by what it holds. One built by another function and passed in is keyed by what that function returns. Nothing needs an annotation.
 
-One program uses one wait key type. A turn has one wait, the key is how the program says what it is waiting for, and the keys leave every backend through one set of helpers built from that type. A program that keys one wait one way and another wait another way is refused. The fix is to name both kinds as constructors of one type. A wait whose key type nothing settles is refused for the same reason, and the fix is to bind the map to a name with its type written down. `aver check` reads one file at a time, so a program whose modules disagree about the key is refused at the compile door, which reads the entry module and the modules it depends on together. `--target wasip2` keys a wait by `Int`. It carries the wait through canonical ABI imports instead of the helpers an external host walks, and a wait keyed by another type is refused for that target.
+One program uses one wait key type. A turn has one wait, the key is how the program says what it is waiting for, and the keys leave every backend through one set of helpers built from that type. The generated loop is the exception. It keys its own wait by `Int`, so in a program that answers a capability of its own, every other wait keyed by another type is carried through an `Int`-keyed wait. That covers the entry and its dependencies alike. The compiler generates one set of helpers per key type, in the `__` namespace of the module that waits, for example `__waitPollByWatch`. They number the keys in the order the map puts them in, wait on the numbered set, and answer the keys the ready numbers stand for, in the same order. The program gets the answer it would have got, and every backend, `--target wasip2` included, sees `Int` keys. A recording holds the `Int`-keyed wait. Outside such a program, a program that keys one wait one way and another wait another way is refused. The fix is to name both kinds as constructors of one type. A wait whose key type nothing settles is refused for the same reason, and the fix is to bind the map to a name with its type written down. `aver check` reads one file at a time, so a program whose modules disagree about the key is refused at the compile door, which reads the entry module and the modules it depends on together. `--target wasip2` keys a wait by `Int`. It carries the wait through canonical ABI imports instead of the helpers an external host walks, and a wait keyed by another type is refused for that target.
 
 A wait set written empty at the call is a turn with nothing to watch. It holds nothing, so it names no key of its own and is read as keyed by `Int`, which is what such a call has always meant. In a program that keys its waits by a type of its own, that reads as a second key type and the program is refused by name. Write the type on that one set:
 

src/capability/work.rs

@@ -1654,7 +1654,7 @@ fn wait_set_key(annotation: &str) -> Option<Type> {
 ///
 /// The generic contract's own `Map<K, Wait.Item>` is the signature rather
 /// than a choice a program made, so a type variable is not an answer.
-fn wait_set_key_of_type(ty: &Type) -> Option<Type> {
+pub(crate) fn wait_set_key_of_type(ty: &Type) -> Option<Type> {
     fn walk(ty: &Type, out: &mut Option<Type>) {
         match ty {
             Type::Map(key, value) if is_wait_item(value) => *out = Some((**key).clone()),

src/ir/pipeline.rs

@@ -1074,30 +1074,48 @@ pub fn front(items: &mut Vec<TopLevel>, cfg: FrontConfig<'_, '_>) -> FrontResult
                 ..phase_one
             },
         }
-    } else if crate::ir::nested_patterns::has_nested_patterns(items) {
-        // Checked as written first, so exhaustiveness, redundancy, type
-        // and shadowing errors name the patterns the user wrote; then the
-        // nested patterns are compiled to flat matches and the lowered
-        // program is checked again, which stamps the nodes the
-        // compilation made.
-        let phase_one = typecheck_gate(items, mode, &items[..user_program_len]);
-        if !phase_one.errors.is_empty() {
-            phase_one
-        } else {
-            let errors = crate::ir::nested_patterns::lower_nested_patterns(
-                items,
-                &phase_one.pattern_ctor_families,
-            );
+    } else if crate::ir::nested_patterns::has_nested_patterns(items)
+        || (!marked.is_empty() && crate::yield_lowering::calls_wait_poll(items))
+    {
+        // A module with no process of its own. Checked as written first, so
+        // exhaustiveness, redundancy, type and shadowing errors name the
+        // patterns the user wrote. Its nested patterns are then compiled to
+        // flat matches and the lowered module is checked again, which stamps
+        // the nodes the compilation made. Last, in a module of a program that
+        // answers a capability of its own, its waits keyed by another type
+        // than `Int` are carried through an `Int`-keyed one, the same way the
+        // entry's are, so they can meet the generated loop's wait in one
+        // program. The keys are read off the lowered module as the last check
+        // stamped it, and what the carrying wrote is checked once more.
+        let carries = !marked.is_empty() && crate::yield_lowering::calls_wait_poll(items);
+        let mut tc = typecheck_gate(items, mode, &items[..user_program_len]);
+        if tc.errors.is_empty() && crate::ir::nested_patterns::has_nested_patterns(items) {
+            let errors =
+                crate::ir::nested_patterns::lower_nested_patterns(items, &tc.pattern_ctor_families);
             fire(PipelineStage::PatternLower, items);
-            if errors.is_empty() {
+            tc = if errors.is_empty() {
                 typecheck(items, mode)
             } else {
-                TypeCheckResult {
-                    errors,
-                    ..phase_one
+                TypeCheckResult { errors, ..tc }
+            };
+        }
+        if tc.errors.is_empty() && carries {
+            let stamped = items.clone();
+            match crate::yield_lowering::carry_waits(items, &stamped, &tc.type_spellings) {
+                Ok(Some(source)) => {
+                    if std::env::var_os("AVER_YIELD_DUMP").is_some() {
+                        eprintln!("{source}");
+                    }
+                    if run_tco {
+                        tco(items);
+                    }
+                    tc = typecheck(items, mode);
                 }
+                Ok(None) => {}
+                Err(errors) => tc = TypeCheckResult { errors, ..tc },
             }
         }
+        tc
     } else {
         typecheck_gate(items, mode, &items[..user_program_len])
     };

src/yield_lowering/carried_waits.rs

@@ -0,0 +1,181 @@
+//! Waits the program keys by a type of its own, carried through an `Int`-keyed
+//! wait (jasisz/aver#1329, process layer v2).
+//!
+//! The generated loop keys its one wait by `Int`: it numbers every item every
+//! parked request waits on. A program that also waits by hand, keyed by a type
+//! of its own (`Map<Watch, Wait.Item>`), would then key waits two ways, and
+//! every backend names the wait's boundary from one key type. So in a program
+//! that answers a capability of its own, each such wait is rewritten to call a
+//! helper generated for its key type, in the reserved `__` namespace of the
+//! module that makes it. The helper numbers the keys in the order the map puts
+//! them in, waits on the numbered set, and hands back the keys the numbers
+//! stand for, in the same order. The program sees the same answer it would
+//! have seen; the backends see one key type.
+//!
+//! No generics: one set of helpers per key type, spelled for that type.
+
+use std::collections::BTreeMap;
+
+use crate::ast::{Expr, FnDef, Spanned, Stmt, TopLevel, Type};
+use crate::codegen::expr_walk::for_each_child_mut;
+
+const WAIT_POLL: &str = "Wait.poll";
+
+/// What was carried: the helpers generated, as source, and their items.
+pub(crate) struct Carried {
+    pub source: String,
+    pub items: Vec<TopLevel>,
+}
+
+fn walk(expr: &mut Spanned<Expr>, visit: &mut impl FnMut(&mut Spanned<Expr>)) {
+    visit(expr);
+    for_each_child_mut(expr, &mut |child| walk(child, visit));
+}
+
+fn walk_fn(fd: &mut FnDef, visit: &mut impl FnMut(&mut Spanned<Expr>)) {
+    let body = std::sync::Arc::make_mut(&mut fd.body);
+    for stmt in body.stmts_mut() {
+        match stmt {
+            Stmt::Binding(_, _, value) | Stmt::Expr(value) => walk(value, visit),
+        }
+    }
+}
+
+fn is_wait_poll(callee: &Expr) -> bool {
+    fn dotted(expr: &Expr) -> Option<String> {
+        match expr {
+            Expr::Ident(name) => Some(name.clone()),
+            Expr::Attr(obj, field) => Some(format!("{}.{field}", dotted(&obj.node)?)),
+            _ => None,
+        }
+    }
+    dotted(callee).as_deref() == Some(WAIT_POLL)
+}
+
+/// Whether a module calls `Wait.poll` anywhere, read off the source alone.
+pub(crate) fn calls_wait_poll(items: &[TopLevel]) -> bool {
+    items.iter().any(|item| {
+        let TopLevel::FnDef(fd) = item else {
+            return false;
+        };
+        fd.body.stmts().iter().any(|stmt| {
+            let value = match stmt {
+                Stmt::Binding(_, _, value) | Stmt::Expr(value) => value,
+            };
+            crate::codegen::expr_walk::any(
+                value,
+                &mut |e| matches!(&e.node, Expr::FnCall(callee, _) if is_wait_poll(&callee.node)),
+            )
+        })
+    })
+}
+
+/// The key of every `Wait.poll` call in `fd`, in walk order, read off the
+/// types the checker stamped on each wait set.
+fn keys_of(fd: &FnDef) -> Vec<Option<Type>> {
+    let mut copy = fd.clone();
+    let mut keys = Vec::new();
+    walk_fn(&mut copy, &mut |expr| {
+        if let Expr::FnCall(callee, args) = &expr.node
+            && is_wait_poll(&callee.node)
+        {
+            keys.push(
+                args.first()
+                    .and_then(|set| set.ty())
+                    .and_then(crate::capability::work::wait_set_key_of_type),
+            );
+        }
+    });
+    keys
+}
+
+/// `Infra.Watch` → `InfraWatch`: the part of a helper's name that says which
+/// key type it carries.
+fn suffix(spelled: &str) -> String {
+    spelled
+        .chars()
+        .filter(|c| c.is_ascii_alphanumeric())
+        .collect()
+}
+
+/// Rewrite every wait of `items` keyed by something other than `Int` to go
+/// through a generated helper for its key. `stamped` is the same module after
+/// a type check, in the same order, so every wait set carries its type. The
+/// functions `skip` names were replaced by generated code and are left alone.
+pub(crate) fn carry(
+    items: &mut Vec<TopLevel>,
+    stamped: &[TopLevel],
+    spellings: &super::TypeSpellings,
+    skip: &dyn Fn(&str) -> bool,
+) -> Result<Option<Carried>, String> {
+    let mut by_fn: BTreeMap<String, Vec<Option<Type>>> = BTreeMap::new();
+    for item in stamped {
+        if let TopLevel::FnDef(fd) = item
+            && !skip(&fd.name)
+        {
+            let keys = keys_of(fd);
+            if keys
+                .iter()
+                .any(|key| key.as_ref().is_some_and(|key| *key != Type::Int))
+            {
+                by_fn.insert(fd.name.clone(), keys);
+            }
+        }
+    }
+    if by_fn.is_empty() {
+        return Ok(None);
+    }
+    // One helper set per key type, named after the key as the module spells it.
+    let mut helpers: BTreeMap<String, String> = BTreeMap::new();
+    for item in items.iter_mut() {
+        let TopLevel::FnDef(fd) = item else { continue };
+        let Some(keys) = by_fn.get(&fd.name) else {
+            continue;
+        };
+        let mut at = 0usize;
+        walk_fn(fd, &mut |expr| {
+            let Expr::FnCall(callee, _) = &mut expr.node else {
+                return;
+            };
+            if !is_wait_poll(&callee.node) {
+                return;
+            }
+            let key = keys.get(at).cloned().flatten();
+            at += 1;
+            let Some(key) = key.filter(|key| *key != Type::Int) else {
+                return;
+            };
+            let spelled = super::spell_type(&key, spellings);
+            let name = format!("__waitPollBy{}", suffix(&spelled));
+            helpers.entry(spelled).or_insert_with(|| name.clone());
+            callee.node = Expr::Ident(name);
+        });
+    }
+    let mut source = String::new();
+    for (key, name) in &helpers {
+        source.push_str(&helper_source(key, name));
+    }
+    let tokens = crate::lexer::Lexer::new(&source)
+        .tokenize()
+        .map_err(|error| error.to_string())?;
+    let parsed = crate::parser::Parser::new_compiler_generated(tokens)
+        .parse()
+        .map_err(|error| error.to_string())?;
+    items.extend(parsed.iter().cloned());
+    Ok(Some(Carried {
+        source,
+        items: parsed,
+    }))
+}
+
+/// The helpers that carry one key type through an `Int`-keyed wait.
+fn helper_source(key: &str, name: &str) -> String {
+    let tail = &name["__waitPollBy".len()..];
+    format!(
+        "\nfn {name}(items: Map<{key}, Wait.Item>, timeoutMs: Int) -> Result<List<{key}>, String>\n    ? \"Waits on a set keyed by {key} through a wait keyed by the position of each key, and answers the keys that were ready, in the order the set puts them in.\"\n    ! [Wait.poll]\n    keys = Map.keys(items)\n    ready = Wait.poll(__waitNumbered{tail}(items, keys, 0, {{}}), timeoutMs)?\n    Result.Ok(__waitKeysAt{tail}(Vector.fromList(keys), ready, []))\n\
+         \nfn __waitNumbered{tail}(items: Map<{key}, Wait.Item>, keys: List<{key}>, next: Int, acc: Map<Int, Wait.Item>) -> Map<Int, Wait.Item>\n    ? \"The same items, keyed by the position of their key in the set's own order.\"\n    match keys\n        [] -> acc\n        [key, ..rest] -> __waitNumbered{tail}(items, rest, next + 1, __waitNumberedAt{tail}(items, key, next, acc))\n\
+         \nfn __waitNumberedAt{tail}(items: Map<{key}, Wait.Item>, key: {key}, next: Int, acc: Map<Int, Wait.Item>) -> Map<Int, Wait.Item>\n    ? \"One item under its position.\"\n    match Map.get(items, key)\n        Option.None -> acc\n        Option.Some(item) -> Map.set(acc, next, item)\n\
+         \nfn __waitKeysAt{tail}(keys: Vector<{key}>, ready: List<Int>, acc: List<{key}>) -> List<{key}>\n    ? \"The keys the ready positions stand for, in the order they were reported.\"\n    match ready\n        [] -> acc\n        [at, ..rest] -> __waitKeysAt{tail}(keys, rest, __waitKeyAt{tail}(keys, at, acc))\n\
+         \nfn __waitKeyAt{tail}(keys: Vector<{key}>, at: Int, acc: List<{key}>) -> List<{key}>\n    ? \"One ready position, as its key.\"\n    match Vector.get(keys, at)\n        Option.None -> acc\n        Option.Some(key) -> List.concat(acc, [key])\n"
+    )
+}

src/yield_lowering/mod.rs

@@ -56,6 +56,7 @@ pub(crate) fn spell_type(ty: &crate::ast::Type, spellings: &TypeSpellings) -> St
 }
 
 mod build;
+mod carried_waits;
 mod coordinator;
 mod lower;
 mod trace;
@@ -210,6 +211,31 @@ impl YieldLoweringReport {
 pub const YIELD_EFFECT: &str = "yield";
 
 /// Whether a function body calls `Run.all()`, which runs the generated loop.
+/// Whether a module calls `Wait.poll` anywhere, read off its source.
+pub fn calls_wait_poll(items: &[TopLevel]) -> bool {
+    carried_waits::calls_wait_poll(items)
+}
+
+/// Carry the waits of a module with no process through an `Int`-keyed wait;
+/// see `carried_waits`. `stamped` is the module after a type check. Answers
+/// the generated source when anything was carried.
+pub fn carry_waits(
+    items: &mut Vec<TopLevel>,
+    stamped: &[TopLevel],
+    spellings: &TypeSpellings,
+) -> Result<Option<String>, Vec<TypeError>> {
+    carried_waits::carry(items, stamped, spellings, &|_| false)
+        .map(|carried| carried.map(|carried| carried.source))
+        .map_err(|parse| {
+            vec![error_at(
+                1,
+                format!(
+                    "internal error carrying this module's waits through an Int-keyed wait: {parse}; please report this program"
+                ),
+            )]
+        })
+}
+
 pub fn calls_run_all(fd: &FnDef) -> bool {
     coordinator::calls_run_all(fd)
 }
@@ -610,6 +636,31 @@ pub fn lower(
         }
     }
 
+    // The generated loop keys its wait by `Int`. A program that answers a
+    // capability of its own carries every other wait it writes through an
+    // `Int`-keyed one, so the loop's wait and the program's own can meet in
+    // one program.
+    if !marked.is_empty() {
+        let carried = carried_waits::carry(items, stamped, type_spellings, &|name| {
+            yield_fns.contains(name)
+        })
+        .map_err(|parse| {
+            vec![error_at(
+                1,
+                format!(
+                    "internal error carrying this module's waits through an Int-keyed wait: {parse}; please report this program"
+                ),
+            )]
+        })?;
+        if let Some(carried) = carried {
+            report.generated.extend(carried.items);
+            match report.loop_source.as_mut() {
+                Some(source) => source.push_str(&carried.source),
+                None => report.loop_source = Some(carried.source),
+            }
+        }
+    }
+
     // Materialize the default export surface before replacing source process
     // names with reserved protocol names (which the underscore rule hides).
     let exports = crate::visibility::collect_module_exports(stamped);

tests/fixtures/run_wait_own_key/aver.toml

@@ -0,0 +1,6 @@
+[providers]
+schema = 1
+
+[[providers.bindings]]
+capability = "Validation"
+work = "Scoring.score"

tests/fixtures/run_wait_own_key/clock.av

@@ -0,0 +1,10 @@
+module Clock
+    kind = capability
+    semantics = effectful
+    intent = "The tick the processes count."
+    exposes [tick]
+
+operation tick() -> Int
+    ? "The number of the next tick."
+    oracle = generative
+    replay = recorded