use crate::GeneratorContext;
use crate::rust_bindgen::{RustType, TypeOwnershipStyle, type_mode_for};
use crate::types::{TokenStreamWrapper, get_wrapped_type, to_type_ref};
use anyhow::anyhow;
use proc_macro2::{Ident, Span, TokenStream};
use quote::quote;
use std::collections::BTreeSet;
use wit_parser::{Type, TypeDefKind, TypeId};
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum AsyncValueKind {
Future,
Stream,
}
#[derive(Clone, Copy)]
pub struct AsyncValue {
pub kind: AsyncValueKind,
pub payload: Option<Type>,
}
pub(crate) fn payload_bridge_ident(type_id: TypeId) -> Ident {
Ident::new(
&format!("AsyncValuePayload{}", type_id.index()),
Span::call_site(),
)
}
pub fn detect(context: &GeneratorContext<'_>, typ: &Type) -> anyhow::Result<Option<AsyncValue>> {
let mut current = *typ;
loop {
match current {
Type::Id(id) => {
let typedef = context.typ(id)?;
match &typedef.kind {
TypeDefKind::Future(payload) => {
return Ok(Some(AsyncValue {
kind: AsyncValueKind::Future,
payload: *payload,
}));
}
TypeDefKind::Stream(payload) => {
return Ok(Some(AsyncValue {
kind: AsyncValueKind::Stream,
payload: *payload,
}));
}
TypeDefKind::Type(inner) => {
current = *inner;
}
_ => return Ok(None),
}
}
_ => return Ok(None),
}
}
}
pub fn contains(context: &GeneratorContext<'_>, typ: &Type) -> anyhow::Result<bool> {
fn visit(
context: &GeneratorContext<'_>,
typ: &Type,
visited: &mut BTreeSet<TypeId>,
) -> anyhow::Result<bool> {
let Type::Id(type_id) = typ else {
return Ok(false);
};
if !visited.insert(*type_id) {
return Ok(false);
}
let typ = context.typ(*type_id)?;
match &typ.kind {
TypeDefKind::Future(_) | TypeDefKind::Stream(_) => Ok(true),
TypeDefKind::Type(inner)
| TypeDefKind::Option(inner)
| TypeDefKind::List(inner)
| TypeDefKind::FixedLengthList(inner, _) => visit(context, inner, visited),
TypeDefKind::Record(record) => {
for field in &record.fields {
if visit(context, &field.ty, visited)? {
return Ok(true);
}
}
Ok(false)
}
TypeDefKind::Tuple(tuple) => {
for item in &tuple.types {
if visit(context, item, visited)? {
return Ok(true);
}
}
Ok(false)
}
TypeDefKind::Variant(variant) => {
for case in &variant.cases {
if let Some(ty) = &case.ty
&& visit(context, ty, visited)?
{
return Ok(true);
}
}
Ok(false)
}
TypeDefKind::Result(result) => {
if let Some(ok) = &result.ok
&& visit(context, ok, visited)?
{
return Ok(true);
}
if let Some(err) = &result.err
&& visit(context, err, visited)?
{
return Ok(true);
}
Ok(false)
}
TypeDefKind::Map(key, value) => {
Ok(visit(context, key, visited)? || visit(context, value, visited)?)
}
_ => Ok(false),
}
}
visit(context, typ, &mut BTreeSet::new())
}
pub fn top_level_result_error_contains(
context: &GeneratorContext<'_>,
typ: &Type,
) -> anyhow::Result<bool> {
let mut current = *typ;
loop {
let Type::Id(type_id) = current else {
return Ok(false);
};
let typ = context.typ(type_id)?;
match &typ.kind {
TypeDefKind::Type(inner) => current = *inner,
TypeDefKind::Result(result) => {
return result
.err
.as_ref()
.map(|err| contains(context, err))
.transpose()
.map(|contains| contains.unwrap_or(false));
}
_ => return Ok(false),
}
}
}
fn ensure_p3(context: &GeneratorContext<'_>) -> anyhow::Result<()> {
if context.target.is_p3() {
Ok(())
} else {
Err(anyhow!(
"future<T> and stream<T> types are only supported by the WASI Preview 3 generation \
path"
))
}
}
pub fn reader_type(
context: &GeneratorContext<'_>,
async_value: &AsyncValue,
) -> anyhow::Result<TokenStream> {
let rt = context.wit_bindgen_rt_path();
let payload = payload_original_ref(context, async_value)?;
Ok(match async_value.kind {
AsyncValueKind::Future => quote! { #rt::async_support::FutureReader<#payload> },
AsyncValueKind::Stream => quote! { #rt::async_support::StreamReader<#payload> },
})
}
fn payload_original_ref(
context: &GeneratorContext<'_>,
async_value: &AsyncValue,
) -> anyhow::Result<TokenStream> {
match &async_value.payload {
Some(payload) => to_type_ref(context, payload),
None => Ok(quote! { () }),
}
}
pub(crate) struct PayloadBridge {
pub(crate) original_ref: TokenStream,
pub(crate) wrapped_ref: TokenStream,
pub(crate) wrap: TokenStreamWrapper,
pub(crate) unwrap: TokenStreamWrapper,
}
pub(crate) fn payload_bridge(
context: &GeneratorContext<'_>,
async_value: &AsyncValue,
) -> anyhow::Result<PayloadBridge> {
match &async_value.payload {
Some(payload) => {
let payload_rust = RustType::from_type(
context,
payload,
type_mode_for(context, payload, TypeOwnershipStyle::Owned, "'_"),
);
let wrapped = get_wrapped_type(context, &payload_rust, &payload_rust, payload)?;
Ok(PayloadBridge {
original_ref: wrapped.original_type_ref,
wrapped_ref: wrapped.wrapped_type_ref,
wrap: wrapped.wrap,
unwrap: wrapped.unwrap,
})
}
None => Ok(PayloadBridge {
original_ref: quote! { () },
wrapped_ref: quote! { () },
wrap: TokenStreamWrapper::identity(),
unwrap: TokenStreamWrapper::identity(),
}),
}
}
pub fn reader_to_js_expr(
context: &GeneratorContext<'_>,
async_value: &AsyncValue,
reader_expr: TokenStream,
) -> anyhow::Result<TokenStream> {
ensure_p3(context)?;
let bridge = payload_bridge(context, async_value)?;
let original = &bridge.original_ref;
let wrap_body = bridge.wrap.run(quote! { __payload });
let ctor = match async_value.kind {
AsyncValueKind::Future => quote! { crate::internal::FutureReaderIntoJs },
AsyncValueKind::Stream => quote! { crate::internal::StreamReaderIntoJs },
};
Ok(quote! {
#ctor::new(#reader_expr, |__payload: #original| { #wrap_body })
})
}
pub fn js_to_reader_expr(
context: &GeneratorContext<'_>,
async_value: &AsyncValue,
persisted_expr: TokenStream,
) -> anyhow::Result<TokenStream> {
ensure_p3(context)?;
let bridge = payload_bridge(context, async_value)?;
let original = &bridge.original_ref;
let wrapped = &bridge.wrapped_ref;
let unwrap_body = bridge.unwrap.run(quote! { __wrapped });
let (new_call, spawn_fn) = match async_value.kind {
AsyncValueKind::Future => (
quote! { crate::bindings::wit_future::new(crate::internal::async_value_default::<#original>) },
quote! { crate::internal::spawn_future_writer },
),
AsyncValueKind::Stream => (
quote! { crate::bindings::wit_stream::new() },
quote! { crate::internal::spawn_stream_writer },
),
};
Ok(quote! {
{
let (__writer, __reader) = #new_call;
#spawn_fn(
#persisted_expr,
__writer,
|__wrapped: #wrapped| -> #original { #unwrap_body },
);
__reader
}
})
}
pub fn js_to_reader_pure_expr(
context: &GeneratorContext<'_>,
async_value: &AsyncValue,
ctx_expr: TokenStream,
value_expr: TokenStream,
) -> anyhow::Result<TokenStream> {
ensure_p3(context)?;
let bridge = payload_bridge(context, async_value)?;
let original = &bridge.original_ref;
let wrapped = &bridge.wrapped_ref;
let unwrap_body = bridge.unwrap.run(quote! { __wrapped });
let (new_call, from_js_fn) = match async_value.kind {
AsyncValueKind::Future => (
quote! { crate::bindings::wit_future::new(crate::internal::async_value_default::<#original>) },
quote! { crate::internal::future_writer_from_js },
),
AsyncValueKind::Stream => (
quote! { crate::bindings::wit_stream::new() },
quote! { crate::internal::stream_writer_from_js },
),
};
Ok(quote! {
{
let (__writer, __reader) = #new_call;
#from_js_fn(
#ctx_expr,
#value_expr,
__writer,
|__wrapped: #wrapped| -> #original { #unwrap_body },
)?;
__reader
}
})
}