obeli_sk_utils/
wit.rs

1use crate::wasm_tools::ExIm;
2use concepts::{
3    FnName, FunctionExtension, FunctionMetadata, IfcFqnName, PackageExtension, PackageIfcFns,
4    PkgFqn,
5};
6use const_format::formatcp;
7use hashbrown::HashMap;
8use id_arena::Arena;
9use indexmap::IndexMap;
10use semver::{BuildMetadata, Prerelease, Version};
11use std::{fmt::Display, path::PathBuf};
12use tracing::{error, warn};
13use wit_component::WitPrinter;
14use wit_parser::{
15    Function, FunctionKind, Handle, Interface, InterfaceId, PackageId, PackageName, Param, Resolve,
16    Span, Stability, Type, TypeDef, TypeDefKind, TypeOwner, UnresolvedPackageGroup, World,
17    WorldItem, WorldKey,
18};
19
20const OBELISK_TYPES_VERSION_MAJOR: u64 = 5;
21const OBELISK_TYPES_VERSION_MINOR: u64 = 0;
22const OBELISK_TYPES_VERSION_PATCH: u64 = 0;
23const OBELISK_TYPES_VERSION: &str = formatcp!(
24    "{OBELISK_TYPES_VERSION_MAJOR}.{OBELISK_TYPES_VERSION_MINOR}.{OBELISK_TYPES_VERSION_PATCH}"
25);
26const OBELISK_TYPES_PACKAGE_NAME: &str = formatcp!("obelisk:types@{OBELISK_TYPES_VERSION}");
27
28pub const WIT_OBELISK_LOG_PACKAGE: [&str; 3] = [
29    "obelisk_log@1.0.0",
30    "obelisk_log@1.0.0.wit",
31    include_str!(concat!(
32        env!("CARGO_MANIFEST_DIR"),
33        "/wit/obelisk_log@1.0.0/obelisk_log@1.0.0.wit"
34    )),
35];
36const WIT_OBELISK_TYPES_PACKAGE_CONTENT: &str = include_str!(concat!(
37    env!("CARGO_MANIFEST_DIR"),
38    "/wit/obelisk_types@5.0.0/obelisk_types@5.0.0.wit"
39));
40pub const WIT_OBELISK_TYPES_PACKAGE: [&str; 3] = [
41    "obelisk_types@5.0.0",
42    "obelisk_types@5.0.0.wit",
43    WIT_OBELISK_TYPES_PACKAGE_CONTENT,
44];
45pub const WIT_OBELISK_WORKFLOW_PACKAGE: [&str; 3] = [
46    "obelisk_workflow@6.0.0",
47    "obelisk_workflow@6.0.0.wit",
48    include_str!(concat!(
49        env!("CARGO_MANIFEST_DIR"),
50        "/wit/obelisk_workflow@6.0.0/obelisk_workflow@6.0.0.wit"
51    )),
52];
53
54pub const WIT_OBELISK_WEBHOOK_PACKAGE: [&str; 3] = [
55    "obelisk_webhook@6.0.0",
56    "obelisk_webhook@6.0.0.wit",
57    include_str!(concat!(
58        env!("CARGO_MANIFEST_DIR"),
59        "/wit/obelisk_webhook@6.0.0/obelisk_webhook@6.0.0.wit"
60    )),
61];
62
63pub(crate) fn wit(resolve: &Resolve, main_package: PackageId) -> Result<String, anyhow::Error> {
64    // print all packages, with the main package as root, others as nested.
65    let ids = packages_except_main(resolve, main_package, false);
66    let mut printer = WitPrinter::new(OutputToString::default());
67    printer.print(resolve, main_package, &ids)?;
68    let wit = printer.output.to_string();
69    Ok(wit)
70}
71
72#[expect(
73    clippy::needless_pass_by_value,
74    reason = "original resolve is consumed; the enriched resolve is returned"
75)]
76pub(crate) fn rebuild_resolve(
77    exim: &ExIm,
78    resolve: Resolve,
79    main_package: PackageId,
80) -> Result<(Resolve, PackageId), anyhow::Error> {
81    let ids = packages_except_main(&resolve, main_package, true);
82    let mut printer = WitPrinter::new(OutputToString::default());
83    printer.print(&resolve, main_package, &ids)?;
84    let wit = printer.output.to_string();
85
86    let (mut resolve, main_pkg_id) = {
87        let wit = replace_obelisk_types(&wit);
88        let group = UnresolvedPackageGroup::parse(PathBuf::new(), &wit)?;
89        let mut resolve = Resolve::new();
90        let main_pkg_id = resolve.push_group(group)?;
91        (resolve, main_pkg_id)
92    };
93    let world_id = resolve
94        .select_world(&[main_pkg_id], None)
95        .expect("default world must be found");
96    let added_interfaces = add_extended_interfaces(exim.get_exports_hierarchy_ext(), &mut resolve)?;
97    resolve
98        .worlds
99        .get_mut(world_id)
100        .expect("id belongs to this resolve")
101        .exports
102        .extend(added_interfaces.into_iter().map(|ifc_id| {
103            (
104                WorldKey::Interface(ifc_id),
105                WorldItem::Interface {
106                    id: ifc_id,
107                    stability: Stability::Unknown,
108                    span: Span::default(),
109                },
110            )
111        }));
112
113    Ok((resolve, main_pkg_id))
114}
115
116fn add_extended_interfaces(
117    exports_hierarchy_ext: &[PackageIfcFns],
118    resolve: &mut Resolve,
119) -> Result<Vec<InterfaceId>, semver::Error> {
120    let mut added_interfaces = Vec::new();
121    // Find necessary handles
122    // Get obelisk:types
123    let obelisk_types_package_name = PackageName {
124        namespace: "obelisk".to_string(),
125        name: "types".to_string(),
126        version: Some(Version {
127            major: OBELISK_TYPES_VERSION_MAJOR,
128            minor: OBELISK_TYPES_VERSION_MINOR,
129            patch: OBELISK_TYPES_VERSION_PATCH,
130            pre: Prerelease::EMPTY,
131            build: BuildMetadata::EMPTY,
132        }),
133    };
134    let obelisk_types_pkg_id =
135        if let Some(id) = resolve.package_names.get(&obelisk_types_package_name) {
136            *id
137        } else {
138            let pkg = wit_parser::Package {
139                name: obelisk_types_package_name,
140                docs: wit_parser::Docs::default(),
141                interfaces: IndexMap::default(),
142                worlds: IndexMap::default(),
143            };
144            let package_name = pkg.name.clone();
145            let ext_pkg_id = resolve.packages.alloc(pkg);
146            resolve.package_names.insert(package_name, ext_pkg_id);
147            ext_pkg_id
148        };
149    // Get obelisk:types/time@VERSION
150    let time_ifc_id = *resolve.packages[obelisk_types_pkg_id]
151        .interfaces
152        .get("time")
153        .expect("`time` interface was added");
154    let time_ifc = &resolve.interfaces[time_ifc_id];
155
156    let (execution_ifc_id, execution_ifc) = find_interface(
157        &IfcFqnName::from_parts("obelisk", "types", "execution", Some(OBELISK_TYPES_VERSION)),
158        resolve,
159        &resolve.interfaces,
160    )
161    .expect(formatcp!(
162        "{OBELISK_TYPES_PACKAGE_NAME} with interface `execution` must be found"
163    ));
164
165    let (join_set_ifc_id, join_set_ifc) = find_interface(
166        &IfcFqnName::from_parts("obelisk", "types", "join-set", Some(OBELISK_TYPES_VERSION)),
167        resolve,
168        &resolve.interfaces,
169    )
170    .expect(formatcp!(
171        "{OBELISK_TYPES_PACKAGE_NAME} with interface `join-set` must be found"
172    ));
173
174    // obelisk:types/execution@VERSION.{execution-id}
175    let type_id_execution_id = {
176        let actual_type_id = *execution_ifc
177            .types
178            .get("execution-id")
179            .expect("`execution-id` must exist");
180        // Create a reference to the type.
181        resolve.types.alloc(TypeDef {
182            name: None,
183            kind: TypeDefKind::Type(Type::Id(actual_type_id)),
184            owner: TypeOwner::Interface(execution_ifc_id),
185            docs: wit_parser::Docs::default(),
186            stability: wit_parser::Stability::default(),
187            span: Span::default(),
188        })
189    };
190    // obelisk:types/execution@VERSION.{join-set}
191    let (type_id_join_set_id, type_id_join_set_id_borrow_handle) = {
192        let actual_type_id = *join_set_ifc
193            .types
194            .get("join-set")
195            .expect("`join-set` must exist");
196        // Create a reference to the type.
197        let type_id_join_set_id = resolve.types.alloc(TypeDef {
198            name: Some("join-set".to_string()),
199            kind: TypeDefKind::Type(Type::Id(actual_type_id)),
200            owner: TypeOwner::Interface(join_set_ifc_id),
201            docs: wit_parser::Docs::default(),
202            stability: wit_parser::Stability::default(),
203            span: Span::default(),
204        });
205        // Create a Handle::Borrow to the reference.
206        let type_id_join_set_id_borrow_handle = resolve.types.alloc(TypeDef {
207            name: None,
208            kind: TypeDefKind::Handle(Handle::Borrow(type_id_join_set_id)),
209            owner: TypeOwner::Interface(join_set_ifc_id),
210            docs: wit_parser::Docs::default(),
211            stability: wit_parser::Stability::default(),
212            span: Span::default(),
213        });
214        (type_id_join_set_id, type_id_join_set_id_borrow_handle)
215    };
216    // obelisk:types/execution.{await-next-extension-error}
217    let type_id_await_next_extension_error = {
218        let actual_type_id = *execution_ifc
219            .types
220            .get("await-next-extension-error")
221            .expect("`await-next-extension-error` must exist");
222        // Create a reference to the type.
223        resolve.types.alloc(TypeDef {
224            name: None,
225            kind: TypeDefKind::Type(Type::Id(actual_type_id)),
226            owner: TypeOwner::Interface(execution_ifc_id),
227            docs: wit_parser::Docs::default(),
228            stability: wit_parser::Stability::default(),
229            span: Span::default(),
230        })
231    };
232    // obelisk:types/execution.{get-extension-error}
233    let type_id_get_extension_error = {
234        let actual_type_id = *execution_ifc
235            .types
236            .get("get-extension-error")
237            .expect("`get-extension-error` must exist");
238        // Create a reference to the type.
239        resolve.types.alloc(TypeDef {
240            name: None,
241            kind: TypeDefKind::Type(Type::Id(actual_type_id)),
242            owner: TypeOwner::Interface(execution_ifc_id),
243            docs: wit_parser::Docs::default(),
244            stability: wit_parser::Stability::default(),
245            span: Span::default(),
246        })
247    };
248    // obelisk:types/execution.{stub-error}
249    let type_id_stub_error = {
250        let actual_type_id = *execution_ifc
251            .types
252            .get("stub-error")
253            .expect("`stub-error` must exist");
254        // Create a reference to the type.
255        resolve.types.alloc(TypeDef {
256            name: None,
257            kind: TypeDefKind::Type(Type::Id(actual_type_id)),
258            owner: TypeOwner::Interface(execution_ifc_id),
259            docs: wit_parser::Docs::default(),
260            stability: wit_parser::Stability::default(),
261            span: Span::default(),
262        })
263    };
264    let type_id_await_next_err_part = type_id_await_next_extension_error;
265    // obelisk:types/time.{schedule-at}
266    let type_id_schedule_at = {
267        let actual_type_id = *time_ifc
268            .types
269            .get("schedule-at")
270            .expect("`schedule-at` must exist");
271        // Create a reference to the type.
272        resolve.types.alloc(TypeDef {
273            name: None,
274            kind: TypeDefKind::Type(Type::Id(actual_type_id)),
275            owner: TypeOwner::Interface(execution_ifc_id),
276            docs: wit_parser::Docs::default(),
277            stability: wit_parser::Stability::default(),
278            span: Span::default(),
279        })
280    };
281
282    for (pkg_fqn, ifc_to_fns) in get_ext_pkg_to_ifc_to_details_map(exports_hierarchy_ext) {
283        let (orig_pkg_fqn, pkg_ext) = pkg_fqn
284            .split_ext()
285            .expect("`get_pkg_to_ifc_to_details_map` filtered by ext packages");
286        let pkg_id = get_or_create_package(pkg_fqn, resolve)?;
287        let (orig_pkg_id, _) = resolve
288            .packages
289            .iter()
290            .find(|(_, found_pkg)| {
291                from_wit_package_name_to_pkg_fqn(&found_pkg.name) == orig_pkg_fqn
292            })
293            .unwrap_or_else(|| {
294                panic!("original package must be not found in resolve: {orig_pkg_fqn}")
295            });
296
297        for (ifc_fqn, fns) in ifc_to_fns {
298            let orig_pkg = resolve.packages.get(orig_pkg_id).expect("id is fresh");
299
300            let orig_ifc_id = *orig_pkg
301                .interfaces
302                .get(ifc_fqn.ifc_name())
303                .unwrap_or_else(|| {
304                    panic!("interface must be found in original resolve: {ifc_fqn}")
305                });
306            let orig_ifc = resolve
307                .interfaces
308                .get(orig_ifc_id)
309                .expect("orig_ifc obtained from orig_resolve");
310
311            let mut types = copy_or_refer_original_types(orig_ifc_id, orig_ifc, &mut resolve.types);
312            match pkg_ext {
313                PackageExtension::ObeliskExt => {
314                    types.insert("execution-id".to_string(), type_id_execution_id);
315                    types.insert("join-set".to_string(), type_id_join_set_id);
316                    types.insert(
317                        "await-next-extension-error".to_string(),
318                        type_id_await_next_extension_error,
319                    );
320                    types.insert(
321                        "get-extension-error".to_string(),
322                        type_id_get_extension_error,
323                    );
324                }
325                PackageExtension::ObeliskSchedule => {
326                    types.insert("execution-id".to_string(), type_id_execution_id);
327                    types.insert("schedule-at".to_string(), type_id_schedule_at);
328                }
329                PackageExtension::ObeliskStub => {
330                    types.insert("execution-id".to_string(), type_id_execution_id);
331                    types.insert("stub-error".to_string(), type_id_stub_error);
332                }
333            }
334
335            let mut ifc = Interface {
336                name: Some(ifc_fqn.ifc_name().to_string()),
337                types,
338                functions: IndexMap::default(),
339                docs: wit_parser::Docs::default(),
340                stability: Stability::default(),
341                package: Some(pkg_id),
342                span: Span::default(),
343                clone_of: None,
344            };
345            for (fn_name, fn_meta) in fns {
346                let (prefix, fn_ext) = fn_meta.split_extension().expect("filtered by ext package");
347                let original_fn = orig_ifc.functions.get(prefix).unwrap_or_else(|| {
348                    panic!("original function {prefix} must be found based on {fn_meta:?}")
349                });
350                let (params, result) = match fn_ext {
351                    FunctionExtension::Submit => {
352                        // -submit: func(join-set: borrow<join-set>, <params>) -> execution-id;
353                        assert_eq!(pkg_ext, PackageExtension::ObeliskExt);
354                        let mut params = vec![Param {
355                            name: generate_param_name("join-set", &original_fn.params),
356                            ty: Type::Id(type_id_join_set_id_borrow_handle),
357                            span: Span::default(),
358                        }];
359                        params.extend_from_slice(&original_fn.params);
360
361                        (params, Some(Type::Id(type_id_execution_id)))
362                    }
363                    FunctionExtension::AwaitNext => {
364                        // -await-next: func(join-set: borrow<join-set>) ->
365                        //  result<return-type, await-next-extension-error>;
366                        // The response id is no longer in the tuple; read it via
367                        // `join-set.last-id`. Structurally identical to `-get`.
368                        assert_eq!(pkg_ext, PackageExtension::ObeliskExt);
369                        let params = vec![Param {
370                            name: "join-set".to_string(),
371                            ty: Type::Id(type_id_join_set_id_borrow_handle),
372                            span: Span::default(),
373                        }];
374                        let result = {
375                            let type_id_result = resolve.types.alloc(TypeDef {
376                                name: None,
377                                kind: TypeDefKind::Result(wit_parser::Result_ {
378                                    ok: original_fn.result,
379                                    err: Some(Type::Id(type_id_await_next_err_part)),
380                                }),
381                                owner: TypeOwner::None,
382                                docs: wit_parser::Docs::default(),
383                                stability: wit_parser::Stability::default(),
384                                span: Span::default(),
385                            });
386                            Some(Type::Id(type_id_result))
387                        };
388                        (params, result)
389                    }
390                    FunctionExtension::Schedule => {
391                        // -schedule: func(schedule-at: schedule-at, <params>) -> execution-id;
392                        assert_eq!(pkg_ext, PackageExtension::ObeliskSchedule);
393                        let schedule_at_param_name =
394                            generate_param_name("schedule-at", &original_fn.params);
395                        let mut params = vec![Param {
396                            name: schedule_at_param_name.clone(),
397                            ty: Type::Id(type_id_schedule_at),
398                            span: Span::default(),
399                        }];
400                        params.extend_from_slice(&original_fn.params);
401                        let result = Some(Type::Id(type_id_execution_id));
402                        (params, result)
403                    }
404                    FunctionExtension::Stub => {
405                        // -stub: func(execution_id: execution-id, original retval) -> result<_, stub-error>;
406                        assert_eq!(pkg_ext, PackageExtension::ObeliskStub);
407                        let mut params = vec![Param {
408                            name: "execution-id".to_string(),
409                            ty: Type::Id(type_id_execution_id),
410                            span: Span::default(),
411                        }];
412                        let Some(return_type) = &original_fn.result else {
413                            unreachable!(
414                                "return types of exported functions are validated in ExImLite"
415                            )
416                        };
417                        let return_type_id = Type::Id(resolve.types.alloc(TypeDef {
418                            name: None,
419                            kind: TypeDefKind::Type(*return_type),
420                            owner: TypeOwner::None,
421                            docs: wit_parser::Docs::default(),
422                            stability: wit_parser::Stability::default(),
423                            span: Span::default(),
424                        }));
425                        params.push(Param {
426                            name: "execution-result".to_string(),
427                            ty: return_type_id,
428                            span: Span::default(),
429                        });
430
431                        let result = {
432                            let type_id_result = resolve.types.alloc(TypeDef {
433                                name: None,
434                                kind: TypeDefKind::Result(wit_parser::Result_ {
435                                    ok: None,
436                                    err: Some(Type::Id(type_id_stub_error)),
437                                }),
438                                owner: TypeOwner::None,
439                                docs: wit_parser::Docs::default(),
440                                stability: wit_parser::Stability::default(),
441                                span: Span::default(),
442                            });
443                            Some(Type::Id(type_id_result))
444                        };
445                        (params, result)
446                    }
447                    FunctionExtension::Get => {
448                        // -get(execution-id) -> result<originalreturn type, get-extension-error>
449                        assert_eq!(pkg_ext, PackageExtension::ObeliskExt);
450                        let params = vec![Param {
451                            name: "execution-id".to_string(),
452                            ty: Type::Id(type_id_execution_id),
453                            span: Span::default(),
454                        }];
455                        let result = Some(Type::Id(resolve.types.alloc(TypeDef {
456                            name: None,
457                            kind: TypeDefKind::Result(wit_parser::Result_ {
458                                ok: original_fn.result,
459                                err: Some(Type::Id(type_id_get_extension_error)),
460                            }),
461                            owner: TypeOwner::None,
462                            docs: wit_parser::Docs::default(),
463                            stability: wit_parser::Stability::default(),
464                            span: Span::default(),
465                        })));
466                        (params, result)
467                    }
468                };
469                let wit_fun = Function {
470                    name: fn_name.to_string(),
471                    kind: FunctionKind::Freestanding,
472                    params,
473                    result,
474                    docs: wit_parser::Docs::default(),
475                    stability: Stability::default(),
476                    span: Span::default(),
477                };
478                ifc.functions.insert(fn_name.to_string(), wit_fun);
479            }
480            // Add Interface to `resolve`.
481            let ifc_id = resolve.interfaces.alloc(ifc);
482            resolve
483                .packages
484                .get_mut(pkg_id)
485                .expect("found or inserted already")
486                .interfaces
487                .insert(ifc_fqn.ifc_name().to_string(), ifc_id);
488
489            // Add the interface
490            added_interfaces.push(ifc_id);
491        }
492    }
493    Ok(added_interfaces)
494}
495
496pub(crate) fn packages_except_main(
497    resolve: &Resolve,
498    main_package: PackageId,
499    sorted: bool,
500) -> Vec<PackageId> {
501    let mut packages = resolve
502        .packages
503        .iter()
504        .map(|(id, _)| id)
505        // The main package would show as a nested package as well
506        .filter(|id| *id != main_package)
507        .collect::<Vec<_>>();
508    if sorted {
509        packages.sort();
510    }
511    packages
512}
513
514// Replace obelisk:types from the actual WASM file because it may not contain all types we are going to need in exported functions.
515fn replace_obelisk_types(wit: &str) -> String {
516    // Replace last character of the first line from ; to {
517    let types_nesting = {
518        let mut types_nesting = WIT_OBELISK_TYPES_PACKAGE_CONTENT.replacen(';', "{", 1);
519        types_nesting.push('}');
520        types_nesting
521    };
522    const TYPES_NESTED_PACKAGE_FIRST_LINE: &str =
523        formatcp!("package {OBELISK_TYPES_PACKAGE_NAME} {{");
524    let wit = remove_nested_package(wit, TYPES_NESTED_PACKAGE_FIRST_LINE);
525    let wit = format!("{wit}\n{types_nesting}");
526    wit
527}
528
529fn generate_param_name(param_name: &str, params: &[Param]) -> String {
530    let orig_param_names: hashbrown::HashSet<&str> = params
531        .iter()
532        .map(|Param { name, .. }| name.as_str())
533        .collect();
534    if orig_param_names.contains(param_name) {
535        for my_char in 'a'..='z' {
536            let name = format!("{param_name}-{my_char}");
537            if !orig_param_names.contains(name.as_str()) {
538                return name;
539            }
540        }
541        warn!("Parameter name `{param_name}` collides with other params {orig_param_names:?}");
542    }
543    param_name.to_string()
544}
545
546fn copy_or_refer_original_types(
547    orig_ifc_id: InterfaceId,
548    orig_ifc: &Interface,
549    resolve_types: &mut Arena<TypeDef>,
550) -> IndexMap<String, id_arena::Id<TypeDef>> {
551    let mut target_types = IndexMap::new();
552    // Copy all imports from original to ext interface. Declared types like records will be referenced instead.
553    for (name, orig_type_id) in &orig_ifc.types {
554        let type_def = resolve_types
555            .get(*orig_type_id)
556            .expect("type def must be found in resolve");
557
558        let allocated_type_id = match type_def.kind {
559            TypeDefKind::Type(_) => resolve_types.alloc(type_def.clone()),
560            _ => {
561                // Create a reference to the type.
562                resolve_types.alloc(TypeDef {
563                    name: None,
564                    kind: TypeDefKind::Type(Type::Id(*orig_type_id)),
565                    owner: TypeOwner::Interface(orig_ifc_id),
566                    docs: wit_parser::Docs::default(),
567                    stability: wit_parser::Stability::default(),
568                    span: Span::default(),
569                })
570            }
571        };
572        target_types.insert(name.clone(), allocated_type_id);
573    }
574    target_types
575}
576
577fn get_or_create_package(
578    pkg_fqn: PkgFqn,
579    resolve: &mut Resolve,
580) -> Result<PackageId, semver::Error> {
581    if let Some((pkg_id, _)) = resolve
582        .packages
583        .iter()
584        .find(|(_, found_pkg)| from_wit_package_name_to_pkg_fqn(&found_pkg.name) == pkg_fqn)
585    {
586        Ok(pkg_id)
587    } else {
588        let pkg = wit_parser::Package {
589            name: from_pkg_fqn_to_wit_package_name(pkg_fqn)?,
590            docs: wit_parser::Docs::default(),
591            interfaces: IndexMap::default(),
592            worlds: IndexMap::default(),
593        };
594        let package_name = pkg.name.clone();
595        let pkg_id = resolve.packages.alloc(pkg);
596        resolve.package_names.insert(package_name, pkg_id);
597        Ok(pkg_id)
598    }
599}
600
601fn get_ext_pkg_to_ifc_to_details_map(
602    exports: &[PackageIfcFns],
603) -> IndexMap<PkgFqn, IndexMap<IfcFqnName, IndexMap<FnName, FunctionMetadata>>> {
604    // Consistent iteration order so that the WIT output is deterministic.
605    // Interfaces are sorted already.
606    let mut pkg_to_ifc_to_details_map: IndexMap<PkgFqn, IndexMap<IfcFqnName, IndexMap<FnName, _>>> =
607        IndexMap::new();
608    for pkg_ifc_fns in exports {
609        if pkg_ifc_fns.ifc_fqn.pkg_fqn_name().is_extension() {
610            let inner_map = pkg_to_ifc_to_details_map
611                .entry(pkg_ifc_fns.ifc_fqn.pkg_fqn_name())
612                .or_default();
613            inner_map.insert(pkg_ifc_fns.ifc_fqn.clone(), pkg_ifc_fns.fns.clone());
614        }
615    }
616    pkg_to_ifc_to_details_map
617}
618
619fn find_interface<'a>(
620    ifc_fqn: &IfcFqnName,
621    resolve: &'_ Resolve,
622    interfaces: &'a Arena<Interface>,
623) -> Option<(InterfaceId, &'a Interface)> {
624    let pkg_id = *resolve
625        .package_names
626        .get(&try_from_ifc_fqn_name(ifc_fqn).ok()?)?;
627    let ifc_id = *resolve.packages[pkg_id]
628        .interfaces
629        .get(ifc_fqn.ifc_name())?;
630    interfaces.get(ifc_id).map(|ifc| (ifc_id, ifc))
631}
632
633fn remove_nested_package(wit_string: &str, nested_package_to_remove: &str) -> String {
634    // Find the start of the namespace
635    let Some(nested_package_start) = wit_string.find(nested_package_to_remove) else {
636        return wit_string.to_string();
637    };
638
639    // Find the opening brace after the namespace
640    let Some(open_brace_index) = wit_string[nested_package_start..]
641        .find('{')
642        .map(|idx| nested_package_start + idx)
643    else {
644        panic!("nested namespace must contain '{{'");
645    };
646
647    // Track brace nesting to find the matching closing brace
648    let mut brace_count = 1;
649    let mut current_index = open_brace_index + 1;
650
651    while current_index < wit_string.len() && brace_count > 0 {
652        match wit_string.chars().nth(current_index) {
653            Some('{') => brace_count += 1,
654            Some('}') => brace_count -= 1,
655            _ => {}
656        }
657        current_index += 1;
658    }
659
660    // If we didn't find the matching brace, return the original string
661    if brace_count != 0 {
662        warn!("Cannot remove the nested package {nested_package_to_remove}");
663        return wit_string.to_string();
664    }
665
666    // Remove the package and its contents
667    format!(
668        "{}\n{}",
669        wit_string[..nested_package_start].trim(),
670        wit_string[current_index..].trim()
671    )
672}
673
674fn try_from_ifc_fqn_name(ifc_fqn: &IfcFqnName) -> Result<PackageName, anyhow::Error> {
675    Ok(PackageName {
676        namespace: ifc_fqn.namespace().to_string(),
677        name: ifc_fqn.package_name().to_string(),
678        version: ifc_fqn
679            .version()
680            .map(semver::Version::parse)
681            .transpose()
682            .inspect_err(|err| {
683                error!(
684                    "cannot parse the version `{:?}` - {err:?}",
685                    ifc_fqn.version()
686                );
687            })?,
688    })
689}
690
691pub(crate) fn from_wit_package_name_to_pkg_fqn(package_name: &PackageName) -> PkgFqn {
692    PkgFqn {
693        namespace: package_name.namespace.clone(),
694        package_name: package_name.name.clone(),
695        version: package_name.version.as_ref().map(ToString::to_string),
696    }
697}
698
699fn from_pkg_fqn_to_wit_package_name(pkg_fqn: PkgFqn) -> Result<PackageName, semver::Error> {
700    Ok(PackageName {
701        namespace: pkg_fqn.namespace,
702        name: pkg_fqn.package_name,
703        version: pkg_fqn
704            .version
705            .as_ref()
706            .map(|v| v.parse())
707            .transpose()
708            .inspect_err(|err| error!("Cannot convert version {:?} - {err:?}", pkg_fqn.version))?,
709    })
710}
711
712/// Build a WIT text map for all exported packages (primary + extension + `obelisk:types`
713/// dependency) from the registry's already-merged export hierarchy.
714pub fn build_wit_deps_map(
715    all_exports: &[PackageIfcFns],
716) -> Result<HashMap<PkgFqn, String /* WIT */>, anyhow::Error> {
717    let mut result = HashMap::new();
718
719    // Group non-extension entries by primary package.
720    let mut primary_pkgs: IndexMap<PkgFqn, Vec<&PackageIfcFns>> = IndexMap::new();
721    for pkg_ifc_fns in all_exports.iter().filter(|e| !e.extension) {
722        primary_pkgs
723            .entry(pkg_ifc_fns.ifc_fqn.pkg_fqn_name())
724            .or_default()
725            .push(pkg_ifc_fns);
726    }
727
728    let has_any_extension = all_exports.iter().any(|e| e.extension);
729
730    for (pkg_fqn, ifc_fns_list) in &primary_pkgs {
731        // Build primary Resolve and print it.
732        let (primary_resolve, primary_pkg_id) = build_primary_resolve(pkg_fqn, ifc_fns_list, None)?;
733        let primary_wit = {
734            let mut printer = WitPrinter::new(OutputToString::default());
735            printer.print(&primary_resolve, primary_pkg_id, &[])?;
736            printer.output.to_string()
737        };
738        result.insert(pkg_fqn.clone(), primary_wit.clone());
739
740        // Collect extension entries whose primary package is pkg_fqn.
741        let ext_for_pkg: Vec<PackageIfcFns> = all_exports
742            .iter()
743            .filter(|e| {
744                e.extension
745                    && e.ifc_fqn
746                        .pkg_fqn_name()
747                        .split_ext()
748                        .map(|(orig, _)| orig == *pkg_fqn)
749                        .unwrap_or(false)
750            })
751            .cloned()
752            .collect();
753
754        if !ext_for_pkg.is_empty() {
755            // Roundtrip through WIT text so that obelisk:types is available in the resolve
756            // (required by add_extended_interfaces to build borrow<join-set> handles etc.).
757            let wit_with_types = replace_obelisk_types(&primary_wit);
758            let group = UnresolvedPackageGroup::parse(PathBuf::new(), &wit_with_types)?;
759            let mut resolve_with_types = Resolve::new();
760            resolve_with_types.push_group(group)?;
761
762            add_extended_interfaces(&ext_for_pkg, &mut resolve_with_types)?;
763
764            // Print each extension package that was just added.
765            for (ext_pkg_id, ext_pkg) in &resolve_with_types.packages {
766                let ext_pkg_fqn = from_wit_package_name_to_pkg_fqn(&ext_pkg.name);
767                if ext_pkg_fqn.is_extension() {
768                    let mut printer = WitPrinter::new(OutputToString::default());
769                    printer.print(&resolve_with_types, ext_pkg_id, &[])?;
770                    result.insert(ext_pkg_fqn, printer.output.to_string());
771                }
772            }
773        }
774    }
775
776    // Always include obelisk:types when extension packages are present.
777    if has_any_extension {
778        let types_pkg_fqn = PkgFqn {
779            namespace: "obelisk".to_string(),
780            package_name: "types".to_string(),
781            version: Some(OBELISK_TYPES_VERSION.to_string()),
782        };
783        result.insert(types_pkg_fqn, WIT_OBELISK_TYPES_PACKAGE_CONTENT.to_string());
784    }
785
786    Ok(result)
787}
788
789/// Build a [`Resolve`] that contains a single primary package with all its interfaces.
790///
791/// Each [`PackageIfcFns`] entry contributes one interface; functions are allocated using
792/// [`crate::wit_builder::allocate_type`] so that named types (record, variant, enum, flags) get
793/// proper `TypeDef` entries.
794pub(crate) fn build_primary_resolve(
795    pkg_fqn: &PkgFqn,
796    ifc_fns_list: &[&PackageIfcFns],
797    world_name: Option<&str>,
798) -> Result<(Resolve, PackageId), anyhow::Error> {
799    use crate::wit_builder::allocate_type;
800
801    let mut resolve = Resolve::new();
802
803    let pkg = wit_parser::Package {
804        name: from_pkg_fqn_to_wit_package_name(pkg_fqn.clone())?,
805        docs: wit_parser::Docs::default(),
806        interfaces: IndexMap::default(),
807        worlds: IndexMap::default(),
808    };
809    let pkg_id = resolve.packages.alloc(pkg);
810    resolve
811        .package_names
812        .insert(resolve.packages[pkg_id].name.clone(), pkg_id);
813
814    // Track already-created interfaces so that multiple PackageIfcFns entries
815    // sharing the same interface name get their functions merged into one interface.
816    let mut ifc_name_to_id: HashMap<String, InterfaceId> = HashMap::new();
817
818    for &pkg_ifc_fns in ifc_fns_list {
819        let ifc_name = pkg_ifc_fns.ifc_fqn.ifc_name().to_string();
820        let ifc_id = if let Some(&existing_id) = ifc_name_to_id.get(&ifc_name) {
821            existing_id
822        } else {
823            let iface = Interface {
824                name: Some(ifc_name.clone()),
825                types: IndexMap::new(),
826                functions: IndexMap::new(),
827                docs: wit_parser::Docs::default(),
828                stability: Stability::default(),
829                package: Some(pkg_id),
830                span: Span::default(),
831                clone_of: None,
832            };
833            let new_id = resolve.interfaces.alloc(iface);
834            resolve
835                .packages
836                .get_mut(pkg_id)
837                .unwrap()
838                .interfaces
839                .insert(ifc_name.clone(), new_id);
840            ifc_name_to_id.insert(ifc_name, new_id);
841            new_id
842        };
843
844        let mut dedup = HashMap::new();
845
846        for (fn_name, fn_metadata) in &pkg_ifc_fns.fns {
847            let wit_params: Vec<Param> = fn_metadata
848                .parameter_types
849                .iter()
850                .map(|p| {
851                    let ty = allocate_type(&mut resolve, ifc_id, &p.type_wrapper, &mut dedup);
852                    Param {
853                        name: p.name.as_ref().to_string(),
854                        ty,
855                        span: Span::default(),
856                    }
857                })
858                .collect();
859
860            let return_tw = fn_metadata.return_type.type_wrapper();
861            let result_type = allocate_type(&mut resolve, ifc_id, &return_tw, &mut dedup);
862
863            let wit_fn = Function {
864                name: fn_name.to_string(),
865                kind: FunctionKind::Freestanding,
866                params: wit_params,
867                result: Some(result_type),
868                docs: wit_parser::Docs::default(),
869                stability: Stability::default(),
870                span: Span::default(),
871            };
872            resolve
873                .interfaces
874                .get_mut(ifc_id)
875                .unwrap()
876                .functions
877                .insert(fn_name.to_string(), wit_fn);
878        }
879
880        // Collect named TypeDefs into the interface's types map.
881        let mut types = IndexMap::new();
882        for (type_id, type_def) in &resolve.types {
883            if type_def.owner == TypeOwner::Interface(ifc_id)
884                && let Some(name) = &type_def.name
885            {
886                types.insert(name.clone(), type_id);
887            }
888        }
889        resolve.interfaces.get_mut(ifc_id).unwrap().types = types;
890    }
891
892    let main_pkg_id = if let Some(world_name) = world_name {
893        let world_exports: IndexMap<WorldKey, WorldItem> = resolve.packages[pkg_id]
894            .interfaces
895            .values()
896            .map(|&ifc_id| {
897                (
898                    WorldKey::Interface(ifc_id),
899                    WorldItem::Interface {
900                        id: ifc_id,
901                        stability: Stability::Unknown,
902                        span: Span::default(),
903                    },
904                )
905            })
906            .collect();
907        // Put the synthesized world in a dedicated `root:component` package to avoid
908        // a package-level dependency cycle.
909        let root_pkg = wit_parser::Package {
910            name: PackageName {
911                namespace: "root".to_string(),
912                name: "component".to_string(),
913                version: None,
914            },
915            docs: wit_parser::Docs::default(),
916            interfaces: IndexMap::default(),
917            worlds: IndexMap::default(),
918        };
919        let root_pkg_id = resolve.packages.alloc(root_pkg);
920        resolve
921            .package_names
922            .insert(resolve.packages[root_pkg_id].name.clone(), root_pkg_id);
923        let world = World {
924            name: world_name.to_string(),
925            docs: wit_parser::Docs::default(),
926            imports: IndexMap::default(),
927            exports: world_exports,
928            package: Some(root_pkg_id),
929            span: Span::default(),
930            includes: vec![],
931            stability: Stability::Unknown,
932        };
933        let world_id = resolve.worlds.alloc(world);
934        resolve
935            .packages
936            .get_mut(root_pkg_id)
937            .unwrap()
938            .worlds
939            .insert(world_name.to_string(), world_id);
940        root_pkg_id
941    } else {
942        pkg_id
943    };
944
945    Ok((resolve, main_pkg_id))
946}
947
948// TODO: Make the wit_component's OutputToString configurable for number of spaces
949#[derive(Default)]
950pub struct OutputToString {
951    indent: usize,
952    output: String,
953    // set to true after newline, then to false after first item is indented.
954    needs_indent: bool,
955}
956
957impl wit_component::Output for OutputToString {
958    fn push_str(&mut self, src: &str) {
959        self.output.push_str(src);
960    }
961
962    fn indent_if_needed(&mut self) -> bool {
963        if self.needs_indent {
964            for _ in 0..self.indent {
965                // Indenting by 4 spaces.
966                self.output.push_str("    ");
967            }
968            self.needs_indent = false;
969            true
970        } else {
971            false
972        }
973    }
974
975    fn indent_start(&mut self) {
976        assert!(
977            !self.needs_indent,
978            "`indent_start` is never called after newline"
979        );
980        self.output.push_str(" {");
981        self.indent += 1;
982        self.newline();
983    }
984
985    fn indent_end(&mut self) {
986        // Note that a `saturating_sub` is used here to prevent a panic
987        // here in the case of invalid code being generated in debug
988        // mode. It's typically easier to debug those issues through
989        // looking at the source code rather than getting a panic.
990        self.indent = self.indent.saturating_sub(1);
991        self.indent_if_needed();
992        self.output.push('}');
993        self.newline();
994    }
995
996    fn newline(&mut self) {
997        self.output.push('\n');
998        self.needs_indent = true;
999    }
1000}
1001impl Display for OutputToString {
1002    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1003        self.output.fmt(f)
1004    }
1005}
1006
1007#[cfg(test)]
1008mod tests {
1009    use crate::{wasm_tools::WasmComponent, wasm_tools::strip_wasm_hash, wit::OutputToString};
1010    use concepts::ComponentType;
1011    use rstest::rstest;
1012    use std::path::PathBuf;
1013    use wit_component::WitPrinter;
1014    use wit_parser::{Resolve, UnresolvedPackageGroup};
1015
1016    #[rstest]
1017    #[case(
1018        test_programs_fibo_workflow_builder::TEST_PROGRAMS_FIBO_WORKFLOW,
1019        ComponentType::Workflow
1020    )]
1021    #[case(
1022        test_programs_fibo_activity_builder::TEST_PROGRAMS_FIBO_ACTIVITY,
1023        ComponentType::Activity
1024    )]
1025    #[case(
1026        test_programs_fibo_webhook_builder::TEST_PROGRAMS_FIBO_WEBHOOK,
1027        ComponentType::WebhookEndpoint
1028    )]
1029    #[case(
1030        test_programs_http_get_activity_builder::TEST_PROGRAMS_HTTP_GET_ACTIVITY,
1031        ComponentType::Activity
1032    )]
1033    #[case(
1034        test_programs_http_get_workflow_builder::TEST_PROGRAMS_HTTP_GET_WORKFLOW,
1035        ComponentType::Workflow
1036    )]
1037    #[case(
1038        test_programs_stub_activity_builder::TEST_PROGRAMS_STUB_ACTIVITY,
1039        ComponentType::ActivityStub
1040    )]
1041    #[case(
1042        test_programs_stub_workflow_builder::TEST_PROGRAMS_STUB_WORKFLOW,
1043        ComponentType::Workflow
1044    )]
1045    #[case(
1046        test_programs_sleep_activity_builder::TEST_PROGRAMS_SLEEP_ACTIVITY,
1047        ComponentType::Workflow
1048    )]
1049    #[case(
1050        test_programs_sleep_workflow_builder::TEST_PROGRAMS_SLEEP_WORKFLOW,
1051        ComponentType::Workflow
1052    )]
1053    fn wit_should_contain_extensions(
1054        #[case] wasm_path: &'static str,
1055        #[case] component_type: ComponentType,
1056    ) {
1057        test_utils::set_up();
1058
1059        let component = WasmComponent::new(wasm_path, component_type).unwrap();
1060        let wasm_path = PathBuf::from(wasm_path);
1061        let wasm_file = strip_wasm_hash(&wasm_path.file_name().unwrap().to_string_lossy());
1062        let wit = component.wit();
1063        // Verify that the generated WIT parses.
1064        let group = UnresolvedPackageGroup::parse(PathBuf::new(), &wit).unwrap();
1065        let mut resolve = Resolve::new();
1066        let main_id = resolve.push_group(group).unwrap();
1067        let ids = resolve
1068            .packages
1069            .iter()
1070            .map(|(id, _)| id)
1071            // The main package would show as a nested package as well
1072            .filter(|id| *id != main_id)
1073            .collect::<Vec<_>>();
1074        let mut printer = WitPrinter::new(OutputToString::default());
1075        printer.print(&resolve, main_id, &ids).unwrap(); // verify it parses
1076        // store original WIT string in snapshots, because that is the `wit()` output.
1077        insta::with_settings!({sort_maps => true, snapshot_suffix => format!("{wasm_file}_wit")}, {insta::assert_snapshot!(wit)});
1078    }
1079}
obeli_sk_utils/wit.rs

obeli_sk_utils/
wit.rs
