asupersync-macros 0.4.6

//! Proc macros for asupersync structured concurrency runtime.
//!
//! This crate provides procedural macros that simplify working with the asupersync
//! async runtime's structured concurrency primitives. The macros handle the boilerplate
//! for creating scopes, spawning tasks, joining results, and racing computations.
//!
//! # Available Macros
//!
//! - [`scope!`] - Create a structured concurrency scope
//! - [`spawn!`] - Spawn a task within the current scope
//! - [`join!`] - Join multiple futures, waiting for all to complete
//! - [`join_all!`] - Join multiple futures into an array
//! - [`race!`] - Race multiple futures, returning the first to complete
//! - [`select!`] - Select over heterogeneous branches with an optional `else` arm
//! - [`#[main]`](macro@main) / [`#[test]`](macro@test) - Production runtime entry attributes
//! - [`session_protocol!`] - Generate typestate session protocols
//! - [`conformance`] - Annotate conformance tests
//! - [`lab_test`] / [`explore_seeds`] - Run deterministic lab-runtime tests
//! - [`ProtoMessage`] / [`ProtoOneof`] - Derive the owned protobuf authoring contract
//!
//! # Contract With `asupersync`
//!
//! The root `asupersync` crate re-exports only the supported runtime DSL:
//! `scope!`, `spawn!`, `join!`, `join_all!`, `race!`, and `select!`, and only
//! when the `proc-macros` feature is enabled.
//!
//! This crate also defines `session_protocol!` and `#[conformance]`, but those
//! remain explicit-path macros on `asupersync_macros`; they are not part of the
//! default root macro contract.
//!
//! # Example
//!
//! ```ignore
//! use asupersync_macros::{scope, spawn, join, race};
//!
//! async fn example(cx: &Cx, state: &mut RuntimeState) {
//!     scope!(cx, state: state, {
//!         let handle1 = spawn!(async { compute_a().await });
//!         let handle2 = spawn!(async { compute_b().await });
//!
//!         // Wait for both
//!         let (result_a, result_b) = join!(handle1, handle2);
//!     });
//! }
//! ```

mod entry;
mod instrument;
mod join;
mod lab_test;
mod proto;
mod race;
mod scope;
mod select;
mod session;
mod spawn;
mod util;

use proc_macro::TokenStream;
use syn::parse_macro_input;

/// Derives the owned `asupersync::grpc::protobuf::ProtoMessage` contract.
///
/// Fields use `#[proto(kind, tag = N)]`, with explicit `optional`,
/// `repeated`, and `packed` modifiers. Maps use
/// `#[proto(map, key = "string", value = "uint64", tag = N)]`; oneofs use an
/// `Option<T>` field annotated with `#[proto(oneof, tags = "N, M")]`.
///
/// Expansion is deterministic and Cargo-only: it invokes no schema compiler
/// and reads no ambient files.
#[proc_macro_derive(ProtoMessage, attributes(proto))]
pub fn derive_proto_message(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as syn::DeriveInput);
    proto::derive_proto_message(&input)
        .unwrap_or_else(syn::Error::into_compile_error)
        .into()
}

/// Derives the owned `asupersync::grpc::protobuf::ProtoOneof` contract.
///
/// Every enum variant must be a one-value tuple variant with its own
/// `#[proto(kind, tag = N)]` attribute.
#[proc_macro_derive(ProtoOneof, attributes(proto))]
pub fn derive_proto_oneof(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as syn::DeriveInput);
    proto::derive_proto_oneof(&input)
        .unwrap_or_else(syn::Error::into_compile_error)
        .into()
}

/// Runs an async `main` function on an asupersync production runtime.
///
/// Supported signatures:
///
/// ```ignore
/// #[asupersync::main]
/// async fn main() {}
///
/// #[asupersync::main(flavor = "current_thread", workers = 1, budget = 128)]
/// async fn main(cx: &asupersync::Cx) -> Result<(), asupersync::Error> {
///     Ok(())
/// }
/// ```
#[proc_macro_attribute]
pub fn main(attr: TokenStream, item: TokenStream) -> TokenStream {
    entry::main_impl(attr, item)
}

/// Runs an async test function on an asupersync production runtime.
///
/// This is distinct from [`#[lab_test]`](macro@lab_test): `#[asupersync::test]`
/// uses the production runtime, while `#[lab_test]` uses deterministic lab
/// runtime seed matrices.
#[proc_macro_attribute]
pub fn test(attr: TokenStream, item: TokenStream) -> TokenStream {
    entry::test_impl(attr, item)
}

/// Creates a structured concurrency scope.
///
/// The `scope!` macro creates an `asupersync::Scope` binding for the
/// current `Cx` region and makes it available as `scope` inside the body.
///
/// Today this is an ergonomic binding helper, not a fresh child-region
/// boundary. For actual child-region ownership and quiescence, call
/// `asupersync::Scope::region` explicitly.
///
/// # Syntax
///
/// ```ignore
/// scope!(cx, {
///     // body with spawned tasks
/// })
/// scope!(cx, state: &mut state, {
///     let _child = spawn!(async { work().await });
/// })
/// ```
///
/// # Arguments
///
/// - `cx` - The capability context (`&Cx`)
/// - `body` - A block containing the scope's work
/// - `state` - Optional runtime state binding used by nested `spawn!` calls
///
/// # Returns
///
/// The result of the scope body.
///
/// # Example
///
/// ```ignore
/// scope!(cx, state: &mut state, {
///     spawn!(async { work_a().await });
///     spawn!(async { work_b().await });
///     // Both tasks are awaited before scope exits
/// })
/// ```
#[proc_macro]
pub fn scope(input: TokenStream) -> TokenStream {
    scope::scope_impl(input)
}

/// Spawns a task within the current scope.
///
/// The `spawn!` macro expands to `asupersync::Scope::spawn_registered`, so it requires
/// ambient `__state` and `__cx` bindings in addition to the target `Scope`.
///
/// The easiest supported path is to use it inside `scope!(..., state: ..., { ... })`.
///
/// # Syntax
///
/// ```ignore
/// spawn!(async { /* work */ })
/// spawn!(async move { /* work with captured values */ })
/// ```
///
/// # Returns
///
/// A `TaskHandle` that can be awaited to get the task's result.
///
/// # Example
///
/// ```ignore
/// let handle = spawn!(async {
///     expensive_computation().await
/// });
/// let result = handle.await;
/// ```
#[proc_macro]
pub fn spawn(input: TokenStream) -> TokenStream {
    spawn::spawn_impl(input)
}

/// Joins multiple futures, waiting for all to complete.
///
/// The `join!` macro polls all branches concurrently inside the enclosing task
/// and returns their outputs as a tuple in input order. A pending branch never
/// blocks ready branches from making progress, so same-duration sleeps complete
/// in one duration rather than the sum of all durations.
///
/// # Syntax
///
/// ```ignore
/// join!(future1, future2, ...)
/// ```
///
/// # Returns
///
/// A tuple of all the futures' results in the order they were specified.
///
/// # Outcome Semantics
///
/// The combined outcome follows the severity lattice:
/// - If all succeed: `Outcome::Ok((r1, r2, ...))`
/// - If any fails: the most severe outcome is propagated
///
/// # Example
///
/// ```ignore
/// let (a, b, c) = join!(
///     fetch_user().await,
///     fetch_profile().await,
///     fetch_settings().await
/// );
/// ```
#[proc_macro]
pub fn join(input: TokenStream) -> TokenStream {
    join::join_impl(input)
}

/// Joins multiple futures into an array, waiting for all to complete.
///
/// The `join_all!` macro is like `join!` but returns an array instead of a
/// tuple. It uses the same concurrent polling expansion, so all branches are
/// driven together within the enclosing task while preserving input order in
/// the returned array.
///
/// # Syntax
///
/// ```ignore
/// join_all!(future1, future2, ...)
/// ```
///
/// # Returns
///
/// An array of all the futures' results in the order they were specified.
/// Since all results must be the same type, this enables easier iteration.
///
/// # Example
///
/// ```ignore
/// let results: [i32; 3] = join_all!(
///     fetch_value(1).await,
///     fetch_value(2).await,
///     fetch_value(3).await
/// );
/// for result in results {
///     println!("{}", result);
/// }
/// ```
#[proc_macro]
pub fn join_all(input: TokenStream) -> TokenStream {
    join::join_all_impl(input)
}

/// Races multiple futures, returning the first to complete — **losers are
/// drained**.
///
/// The `race!` macro expands to the drain-correct
/// `asupersync::Cx::race_drained*` family: each branch is
/// spawned as a region task and resolved through
/// `asupersync::Scope::race_all`, so every losing branch is
/// protocol-cancelled **and drained** (awaited to termination) before the macro
/// returns. This is the drain guarantee that differentiates `race!` from a
/// plain drop-the-losers select.
///
/// Because branches run as spawned tasks, each branch and its output must be
/// `Send + 'static`, and `cx` must be a runtime-wired context carrying spawn
/// authority. For a lower-level drop-on-cancel select over non-`'static`
/// inline futures, call `asupersync::Cx::race` directly.
///
/// # Syntax
///
/// ```ignore
/// race!(cx, { future1, future2, ... })
/// race!(cx, { "name" => future1, "other" => future2, ... })
/// race!(cx, timeout: Duration::from_secs(5), { future1, future2, ... })
/// ```
///
/// # Returns
///
/// The result of the winning future.
///
/// # Loser Cleanup
///
/// All non-winning branches are cancelled and drained: the macro does not
/// return until each loser task has terminated, so obligations and finalizers
/// held by a loser are resolved rather than abandoned. (On the `timeout:` path,
/// an elapsed deadline abandons the whole race by drop, matching
/// `asupersync::Cx::race_drained_timeout`.)
///
/// # Example
///
/// ```ignore
/// let result = race!(cx, {
///     primary_service.fetch().await,
///     backup_service.fetch().await,
/// });
/// // One completed; the loser was cancelled AND drained before this returned.
/// ```
#[proc_macro]
pub fn race(input: TokenStream) -> TokenStream {
    race::race_impl(input)
}

/// Selects over heterogeneous branches — **losers are drained** — with an
/// optional non-blocking `else` arm.
///
/// `select!` is the N-ary, heterogeneous member of the race family. Each branch
/// awaits its own future (the branch types may differ) and runs a handler arm;
/// every handler must yield the same result type `R`. It lifts the fixed
/// `Race2`/`Race3`/`Race4` arity ceiling and is the drain-correct alternative
/// to `tokio::select!`.
///
/// # Two forms
///
/// **Blocking, drain-correct** (no `else` arm): each branch is rewritten into
/// `async move { let <pat> = <future>.await; <handler> }`, and the homogeneous
/// per-branch list routes through
/// `asupersync::Cx::race_drained`. The first branch to win
/// resolves the `select!`; every loser is protocol-cancelled **and drained**
/// (awaited to termination) before the macro returns. Resolves to
/// `Result<R, JoinError>`. Branch futures and `R` must be `Send + 'static`, and
/// `cx` must carry spawn authority (`Cx<cap::All>`).
///
/// **Non-blocking default** (trailing `else => <handler>` arm): each branch is
/// polled **exactly once** in source order; the first ready branch wins,
/// otherwise the `else` handler runs immediately. This is the Go-style
/// `default` arm — it never waits, so it does not drain; not-ready branches are
/// dropped. Resolves to `R`.
///
/// # Determinism / tie-break
///
/// Every `select!` is replay-deterministic: the same seed always produces the
/// same winner. The blocking form resolves through the runtime drain engine
/// (`asupersync::Scope::race_all`), which breaks ties among
/// same-turn-ready branches with the lab's **seeded** scheduler RNG — fixed by
/// the seed, not by source position. The `else` form polls in strict **source
/// order** and takes the first ready branch. The `biased` keyword is accepted
/// on the blocking form for `tokio::select!` familiarity and documents that
/// selection is deterministic; it does not impose strict source order — use the
/// `else` form for that.
///
/// # Syntax
///
/// ```ignore
/// // blocking, drain-correct
/// let r = select!(cx, {
///     a = primary.fetch()  => use_primary(a),
///     b = backup.fetch()   => use_backup(b),
/// })?;
///
/// // explicit source-order tie-break
/// let r = select!(cx, biased, {
///     a = fast()  => a,
///     b = slow()  => b,
/// })?;
///
/// // non-blocking Go-style default
/// let r = select!(cx, {
///     a = try_recv() => a,
///     else => default_value(),
/// });
/// ```
#[proc_macro]
pub fn select(input: TokenStream) -> TokenStream {
    select::select_impl(input)
}

/// Instruments a function or impl method with a tracing span.
///
/// The generated wrapper uses `asupersync::tracing_compat`, so it creates real
/// spans when `tracing-integration` is enabled and becomes a no-op when tracing
/// is disabled.
///
/// Supported arguments:
///
/// - `name = "custom_name"` overrides the span name
/// - `level = "trace" | "debug" | "info" | "warn" | "error"` sets span level
/// - `skip(arg1, arg2, ...)` excludes arguments from captured fields
///
/// # Examples
///
/// ```ignore
/// use asupersync::tracing_compat::instrument;
///
/// #[instrument]
/// async fn load_user(user_id: u64) -> Result<(), Error> {
///     Ok(())
/// }
///
/// #[instrument(name = "cache_refresh", level = "debug", skip(secret))]
/// fn refresh(secret: &Secret, key: &str) {}
/// ```
#[proc_macro_attribute]
pub fn instrument(attr: TokenStream, item: TokenStream) -> TokenStream {
    instrument::instrument_impl(attr, item)
}

/// Marks a test with the specification section and requirement it validates.
///
/// # Syntax
///
/// ```ignore
/// #[conformance(spec = "3.2.1", requirement = "Region close waits for all children")]
/// #[test]
/// fn test_region_close_waits() { /* ... */ }
/// ```
///
/// The macro is validation-only: it checks that `spec` and `requirement` are
/// present and string literals, then leaves the item unchanged.
#[proc_macro_attribute]
pub fn conformance(attr: TokenStream, item: TokenStream) -> TokenStream {
    match parse_conformance_args(&attr) {
        Ok(_) => item,
        Err(message) => util::compile_error(&message).into(),
    }
}

/// Runs a deterministic lab-runtime test with optional seed matrices.
///
/// Supported function shapes:
///
/// ```ignore
/// #[lab_test]
/// fn raw_lab(lab: &mut asupersync::lab::LabRuntime) {
///     // create tasks, advance virtual time, inspect lab state
/// }
///
/// #[lab_test(seeds = 0..16, chaos)]
/// async fn async_body(cx: &asupersync::cx::Cx) {
///     // run under a root lab task with automatic quiescence/oracle checks
/// }
/// ```
#[proc_macro_attribute]
pub fn lab_test(attr: TokenStream, item: TokenStream) -> TokenStream {
    lab_test::lab_test_impl(attr, item)
}

/// Runs a deterministic lab-runtime body across a seed sweep.
///
/// The body is invoked once per seed with a fresh
/// `asupersync::lab::LabRuntime`. The generated test drains each
/// run to quiescence, aggregates trace equivalence-class coverage, and reports
/// failing seeds with replay-friendly reproducer commands.
///
/// Supported arguments:
///
/// - `base = N` or `base_seed = N` sets the first seed
/// - `count = N` sets the number of seeds
/// - `seeds = START..END` uses an exclusive range
/// - `workers = N` or `worker_count = N` sets the lab worker count
/// - `max_steps = N` sets the per-seed step limit
/// - `chaos` enables the light deterministic chaos profile
///
/// # Example
///
/// ```ignore
/// #[explore_seeds(seeds = 0..32, workers = 2)]
/// fn cancellation_matrix(lab: &mut asupersync::lab::LabRuntime) {
///     // build the per-seed scenario; the macro drains and checks it
/// }
/// ```
#[proc_macro_attribute]
pub fn explore_seeds(attr: TokenStream, item: TokenStream) -> TokenStream {
    lab_test::explore_seeds_impl(attr, item)
}

/// Generates typestate-encoded session types from a protocol DSL.
///
/// The macro takes a protocol specification and generates a module containing
/// message structs, paired session type aliases (initiator + responder), and
/// constructor functions. The responder type is the dual of the initiator:
/// `Send`↔`Recv`, `Select`↔`Offer`.
///
/// # Syntax
///
/// ```ignore
/// session_protocol! {
///     module_name<T> for ObligationVariant {
///         msg MessageName;
///         msg MessageWithFields { field: Type };
///
///         send MessageName => select {
///             send T => end,
///             send OtherMsg => end,
///         }
///     }
/// }
/// ```
///
/// # Body Actions
///
/// - `send Type => body` — send a value, then continue
/// - `recv Type => body` — receive a value, then continue
/// - `select { a, b }` — local choice (becomes `Offer` for responder)
/// - `offer { a, b }` — remote choice (becomes `Select` for responder)
/// - `loop { body }` — recursion point (generates `renew_loop` constructor)
/// - `continue` — jump back to enclosing `loop`
/// - `end` — protocol termination
///
/// # Generated Items
///
/// - `pub mod <name>` containing:
///   - Message structs with `Debug, Clone` (+ `Copy` for unit structs)
///   - `InitiatorSession` type alias
///   - `ResponderSession` type alias
///   - `new_session(channel_id) -> (Chan<Initiator, ...>, Chan<Responder, ...>)`
///   - (if `loop` used) `InitiatorLoop`, `ResponderLoop` type aliases
///   - (if `loop` used) `renew_loop(channel_id)` constructor
///
/// # Example
///
/// ```ignore
/// session_protocol! {
///     lease for Lease {
///         msg AcquireMsg;
///         msg RenewMsg;
///         msg ReleaseMsg;
///
///         send AcquireMsg => loop {
///             select {
///                 send RenewMsg => continue,
///                 send ReleaseMsg => end,
///             }
///         }
///     }
/// }
/// ```
#[proc_macro]
pub fn session_protocol(input: TokenStream) -> TokenStream {
    session::session_protocol_impl(input)
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct ConformanceArgs {
    spec: String,
    requirement: String,
}

fn parse_conformance_args(attr: &TokenStream) -> Result<ConformanceArgs, String> {
    parse_conformance_args_str(&attr.to_string())
}

fn parse_conformance_args_str(input: &str) -> Result<ConformanceArgs, String> {
    let raw = input.trim();
    if raw.is_empty() {
        return Err("conformance attribute requires arguments".to_string());
    }

    let mut spec = None;
    let mut requirement = None;

    for part in split_args(raw) {
        let part = part.trim();
        if part.is_empty() {
            continue;
        }
        let (key, value) = split_key_value(part)?;
        let value = parse_string_literal(value)?;
        match key {
            "spec" => spec = Some(value),
            "requirement" => requirement = Some(value),
            other => {
                return Err(format!(
                    "conformance attribute has unknown key '{other}', expected 'spec' or 'requirement'"
                ));
            }
        }
    }

    let spec = spec.ok_or_else(|| "conformance attribute missing 'spec'".to_string())?;
    let requirement =
        requirement.ok_or_else(|| "conformance attribute missing 'requirement'".to_string())?;

    Ok(ConformanceArgs { spec, requirement })
}

fn split_args(input: &str) -> Vec<String> {
    let mut parts = Vec::new();
    let mut current = String::new();
    let mut in_string = false;
    let mut escape = false;

    for ch in input.chars() {
        if in_string {
            current.push(ch);
            if escape {
                escape = false;
                continue;
            }
            if ch == '\\' {
                escape = true;
            } else if ch == '"' {
                in_string = false;
            }
            continue;
        }

        match ch {
            '"' => {
                in_string = true;
                current.push(ch);
            }
            ',' => {
                parts.push(current);
                current = String::new();
            }
            _ => current.push(ch),
        }
    }

    if !current.trim().is_empty() {
        parts.push(current);
    }

    parts
}

fn split_key_value(input: &str) -> Result<(&str, &str), String> {
    let mut iter = input.splitn(2, '=');
    let key = iter
        .next()
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .ok_or_else(|| "conformance attribute expects key = \"value\" pairs".to_string())?;
    let value = iter
        .next()
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .ok_or_else(|| format!("conformance attribute missing value for '{key}'"))?;
    Ok((key, value))
}

fn parse_string_literal(input: &str) -> Result<String, String> {
    let trimmed = input.trim();
    if !trimmed.starts_with('"') || !trimmed.ends_with('"') {
        return Err(format!(
            "conformance attribute values must be string literals, got: {trimmed}"
        ));
    }
    let inner = &trimmed[1..trimmed.len() - 1];
    let mut out = String::new();
    let mut chars = inner.chars();
    while let Some(ch) = chars.next() {
        if ch == '\\' {
            let next = chars.next().ok_or_else(|| {
                "conformance attribute contains dangling escape sequence".to_string()
            })?;
            match next {
                '\\' => out.push('\\'),
                '"' => out.push('"'),
                'n' => out.push('\n'),
                'r' => out.push('\r'),
                't' => out.push('\t'),
                other => {
                    return Err(format!(
                        "conformance attribute contains unsupported escape: \\{other}"
                    ));
                }
            }
        } else {
            out.push(ch);
        }
    }
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::parse_conformance_args_str;

    #[test]
    fn parse_conformance_args_ok() {
        let args =
            parse_conformance_args_str(r#"spec = "3.2.1", requirement = "Region close waits""#)
                .unwrap();
        assert_eq!(args.spec, "3.2.1");
        assert_eq!(args.requirement, "Region close waits");
    }

    #[test]
    fn parse_conformance_args_missing_spec() {
        let err = parse_conformance_args_str(r#"requirement = "Region close waits""#).unwrap_err();
        assert!(err.contains("missing 'spec'"));
    }

    #[test]
    fn parse_conformance_args_missing_requirement() {
        let err = parse_conformance_args_str(r#"spec = "3.2.1""#).unwrap_err();
        assert!(err.contains("missing 'requirement'"));
    }
}