1 use anyhow::anyhow;
2 use clap::Parser;
3 use std::collections::HashMap;
4 use std::ffi::OsStr;
5 use std::fs;
6 use std::path::{Path, PathBuf};
7 use wasmtime::component::{Component, Linker, bindgen};
8 use wasmtime::{Engine, Store};
9 
10 // Generates bindings for the plugin world defined in the wit/calculator.wit file.
11 bindgen!("plugin");
12 
13 /// A CLI that implements a plugin-based calculator whose
14 /// operations are implemented by independent WebAssembly components.
15 
16 #[derive(Parser)]
17 struct BinaryOperation {
18     /// The name of the operation
19     op: String,
20     /// The first operand
21     x: i32,
22     /// The second operand
23     y: i32,
24 }
25 
26 fn lookup_plugin<'a>(
27     state: &'a mut CalculatorState,
28     op: &str,
29 ) -> anyhow::Result<&'a mut PluginDesc> {
30     state
31         .plugin_descs
32         .get_mut(op)
33         .ok_or(anyhow!("unknown operation {}", op))
34 }
35 
36 impl BinaryOperation {
37     fn run(self, state: &mut CalculatorState) -> anyhow::Result<()> {
38         let desc = lookup_plugin(state, self.op.as_ref())?;
39         let result = desc.plugin.call_evaluate(&mut desc.store, self.x, self.y)?;
40         println!("{}({}, {}) = {}", self.op, self.x, self.y, result);
41         Ok(())
42     }
43 }
44 
45 pub struct CalculatorState {
46     // Mapping from operation names to descriptors for the plugins
47     pub plugin_descs: HashMap<String, PluginDesc>,
48 }
49 
50 impl Default for CalculatorState {
51     fn default() -> Self {
52         Self::new()
53     }
54 }
55 
56 impl CalculatorState {
57     pub fn new() -> Self {
58         CalculatorState {
59             plugin_descs: HashMap::new(),
60         }
61     }
62 }
63 
64 pub struct PluginDesc {
65     pub plugin: Plugin,
66     pub store: wasmtime::Store<()>,
67 }
68 
69 fn load_plugin(
70     state: &mut CalculatorState,
71     engine: &Engine,
72     linker: &Linker<()>,
73     path: PathBuf,
74 ) -> wasmtime::Result<()> {
75     println!("Loading plugin from file {:?}", path);
76 
77     // Creates a component from a .wasm file
78     let component = Component::from_file(engine, &path)?;
79 
80     let (plugin_name, plugin, store) = {
81         // Creates a Store for each plugin.
82         // A Store represents dynamic state, so each plugin
83         // has its own Store.
84         let mut store = Store::new(engine, ());
85         // Instantiate the plugin to the store we just created.
86         let plugin = Plugin::instantiate(&mut store, &component, linker)?;
87         (plugin.call_get_plugin_name(&mut store)?, plugin, store)
88     };
89 
90     state
91         .plugin_descs
92         .insert(plugin_name, PluginDesc { plugin, store });
93 
94     Ok(())
95 }
96 
97 fn load_plugins(state: &mut CalculatorState, plugins_dir: &Path) -> wasmtime::Result<()> {
98     let engine = Engine::default();
99     let linker: Linker<()> = Linker::new(&engine);
100 
101     if !plugins_dir.is_dir() {
102         anyhow::bail!("plugins directory does not exist");
103     }
104 
105     for entry in fs::read_dir(plugins_dir)? {
106         let path = entry?.path();
107         if path.is_file() && path.extension().and_then(OsStr::to_str) == Some("wasm") {
108             load_plugin(state, &engine, &linker, path)?;
109         }
110     }
111     Ok(())
112 }
113 
114 #[derive(Parser)]
115 #[command(name = "calculator-host", version = env!("CARGO_PKG_VERSION"))]
116 #[command(about = "A calculator with plugin support")]
117 struct Args {
118     #[command(flatten)]
119     op: BinaryOperation,
120 
121     #[arg(long, help = "Plugin directory")]
122     plugins: PathBuf,
123 }
124 
125 fn main() -> anyhow::Result<()> {
126     // Get plugins directory
127     let args = Args::parse();
128 
129     // Initialize mapping from plugin names to plugins
130     let mut state = CalculatorState::new();
131 
132     // Load plugins from plugins directory
133     load_plugins(&mut state, args.plugins.as_path())?;
134 
135     // Evaluate the expression given on the command line
136     args.op.run(&mut state)
137 }
138