H-86: Add minor hEngine compatibility with TypeScript for behaviors and initialization

apps/sim-engine/lib/execution/Cargo.toml

@@ -21,6 +21,8 @@ rand = "0.8.5"
 rayon = "1.5.3"
 serde = { version = "1.0.138", features = ["derive"] }
 serde_json = "1.0.82"
+swc = "0.261.29"
+swc_common = "0.31.11"
 thiserror = "1.0.31"
 tokio = { version = "1.19.2", features = ["macros", "rt", "sync", "process", "time"] }
 tracing = "0.1.35"

apps/sim-engine/lib/execution/src/package/simulation/init/js_py.rs

@@ -8,7 +8,7 @@ use stateful::{agent::Agent, field::FieldSpecMapAccessor};
 
 pub use self::{
     message::{FailedMessage, JsPyInitTaskMessage, StartMessage, SuccessMessage},
-    task::{JsInitTask, PyInitTask},
+    task::{TsInitTask, JsInitTask, PyInitTask},
 };
 use crate::{
     package::simulation::{
@@ -46,6 +46,9 @@ impl InitPackage for JsPyInit {
             InitialStateName::InitJs => InitTask::JsInitTask(JsInitTask {
                 initial_state_source: self.initial_state.src.clone(),
             }),
+            InitialStateName::InitTs => InitTask::TsInitTask(TsInitTask {
+                initial_state_source: self.initial_state.src.clone(),
+            }),
             name => {
                 // should be unreachable
                 return Err(Error::from(format!(
@@ -89,7 +92,7 @@ impl InitPackageCreator for JsPyInitCreator {
         _accessor: FieldSpecMapAccessor,
     ) -> Result<Box<dyn InitPackage>> {
         match &init_config.initial_state.name {
-            InitialStateName::InitPy | InitialStateName::InitJs => Ok(Box::new(JsPyInit {
+            InitialStateName::InitPy | InitialStateName::InitJs | InitialStateName::InitTs => Ok(Box::new(JsPyInit {
                 initial_state: init_config.initial_state.clone(),
                 comms,
             })),

apps/sim-engine/lib/execution/src/package/simulation/init/js_py/task.rs

@@ -3,7 +3,7 @@ use crate::{
         js_py::{JsPyInitTaskMessage, StartMessage},
         InitTaskMessage,
     },
-    runner::MessageTarget,
+    runner::{Language, MessageTarget},
     task::{TargetedTaskMessage, Task, TaskMessage},
     worker::WorkerHandler,
     Result,
@@ -33,6 +33,34 @@ impl WorkerHandler for JsInitTask {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct TsInitTask {
+    pub initial_state_source: String,
+}
+
+impl Task for TsInitTask {
+    fn name(&self) -> &'static str {
+        "TsInit"
+    }
+}
+
+impl WorkerHandler for TsInitTask {
+    fn start_message(&self) -> Result<TargetedTaskMessage> {
+        let javascript_source =
+            Language::TypeScript.compile_source("init.ts", &self.initial_state_source)?;
+
+        let jspy_init_task_msg = JsPyInitTaskMessage::Start(StartMessage {
+            initial_state_source: javascript_source,
+        });
+
+        let init_task_msg = InitTaskMessage::JsPyInitTaskMessage(jspy_init_task_msg);
+        Ok(TargetedTaskMessage {
+            target: MessageTarget::JavaScript,
+            payload: TaskMessage::Init(init_task_msg),
+        })
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct PyInitTask {
     pub initial_state_source: String,

apps/sim-engine/lib/execution/src/package/simulation/init/state.rs

@@ -6,6 +6,7 @@ pub enum InitialStateName {
     InitJson,
     InitPy,
     InitJs,
+    InitTs,
 }
 
 #[derive(Deserialize, Serialize, Debug, Clone)]

apps/sim-engine/lib/execution/src/package/simulation/init/task.rs

@@ -1,5 +1,5 @@
 use crate::{
-    package::simulation::init::js_py::{JsInitTask, PyInitTask},
+    package::simulation::init::js_py::{TsInitTask, JsInitTask, PyInitTask},
     task::{
         StoreAccessValidator, TargetedTaskMessage, Task, TaskDistributionConfig, TaskSharedStore,
     },
@@ -8,24 +8,28 @@ use crate::{
     Error, Result,
 };
 
+
 /// All init package tasks are registered in this enum
 #[derive(Clone, Debug)]
 pub enum InitTask {
     JsInitTask(JsInitTask),
+    TsInitTask(TsInitTask),
     PyInitTask(PyInitTask),
 }
 
 impl Task for InitTask {
     fn name(&self) -> &'static str {
         match self {
             Self::JsInitTask(task) => task.name(),
+            Self::TsInitTask(task) => task.name(),
             Self::PyInitTask(task) => task.name(),
         }
     }
 
     fn distribution(&self) -> TaskDistributionConfig {
         match self {
             Self::JsInitTask(task) => task.distribution(),
+            Self::TsInitTask(task) => task.distribution(),
             Self::PyInitTask(task) => task.distribution(),
         }
     }
@@ -47,6 +51,7 @@ impl WorkerHandler for InitTask {
     fn start_message(&self) -> Result<TargetedTaskMessage> {
         match self {
             Self::JsInitTask(inner) => inner.start_message(),
+            Self::TsInitTask(inner) => inner.start_message(),
             Self::PyInitTask(inner) => inner.start_message(),
         }
     }

apps/sim-engine/lib/execution/src/package/simulation/state/behavior_execution/behavior.rs

@@ -31,6 +31,19 @@ impl Behavior {
     pub fn language(&self) -> Result<Language> {
         Language::from_file_name(&self.name)
     }
+
+    /// Returns the source code the runner should execute internally, compiling
+    /// the user code if necessary.
+    /// May analyze the source code for validity.
+    pub fn get_runner_source(&self) -> Result<Option<String>> {
+        let language = self.language()?;
+        let compiled = self
+            .behavior_src
+            .as_ref()
+            .map(|src| language.compile_source(&self.name, &src));
+
+        compiled.transpose()
+    }
 }
 
 impl fmt::Debug for Behavior {

apps/sim-engine/lib/execution/src/package/simulation/state/behavior_execution/config.rs

@@ -118,16 +118,19 @@ pub fn exp_init_message(
             let shared = behavior.shared();
             let keys = behavior.keys();
 
-            let language = Language::from_file_name(file_name)
-                .map_err(|_| Error::from("Couldn't get language from behavior file name"))?;
+            let language = shared.language().map_err(|_| {
+                Error::from(format!(
+                    "Couldn't get language from behavior file name {file_name}"
+                ))
+            })?;
             let id = behavior_ids
                 .name_to_index
                 .get(shared.name.as_bytes())
                 .ok_or_else(|| Error::from("Couldn't get index from behavior name"))?;
             let source = shared
-                .behavior_src
-                .clone()
+                .get_runner_source()?
                 .ok_or_else(|| Error::from("SharedBehavior didn't have an attached source"))?;
+
             let required_field_keys = keys
                 .inner
                 .iter()

apps/sim-engine/lib/execution/src/package/simulation/state/behavior_execution/package.js

@@ -11,6 +11,8 @@ const BEHAVIOR_IDS_FIELD_KEY = "_PRIVATE_7_behavior_ids";
 // middle of running the behavior execution package, but `__behaviors`
 // contains behavior ids and shouldn't be modified.
 
+const supported_languages = ["JavaScript", "TypeScript"];
+
 const prepare_user_trace = (error, trace) => {
   let behavior_index = -1;
   for (var i = trace.length - 1; i >= 0; --i) {
@@ -73,14 +75,17 @@ const load_behaviors = (experiment, behavior_descs) => {
   const behaviors = {};
   for (var i = 0; i < behavior_descs.length; ++i) {
     const desc = behavior_descs[i];
-    if (desc.language !== "JavaScript") {
+
+    if (!supported_languages.includes(desc.language)) {
       behaviors[desc.id] = {
         language: desc.language,
       };
+
       continue;
     }
 
     const code = desc.source;
+
     let fn;
     try {
       fn = new Function(
@@ -102,7 +107,6 @@ const load_behaviors = (experiment, behavior_descs) => {
         trace.msg;
       throw new Error(JSON.stringify(trace));
     }
-
     let t = typeof fn;
     if (t !== "function") {
       throw new TypeError(
@@ -209,7 +213,7 @@ export const run_task = (
       // We do this because behavior ids are shallow-loaded and
       // `b_id` is an Arrow Vec rather than a clean array
       const behavior = experiment.behaviors[[b_id.get(0), b_id.get(1)]];
-      if (behavior.language !== "JavaScript") {
+      if (!supported_languages.includes(behavior.language)) {
         // TODO: A simple optimization would be to count the number of
         //       next-up behaviors in each language (other than JS) and
         //       (ignoring ties) choose the language with the most. This

apps/sim-engine/lib/execution/src/package/simulation/task.rs

@@ -2,12 +2,9 @@ use crate::{
     package::simulation::{
         context::ContextTask, init::InitTask, output::OutputTask, state::StateTask,
     },
-    task::{
-        StoreAccessValidator, TargetedTaskMessage, Task, TaskDistributionConfig, TaskMessage,
-        TaskSharedStore,
-    },
+    task::{StoreAccessValidator, Task, TaskDistributionConfig, TaskSharedStore, TargetedTaskMessage, TaskMessage},
     worker::WorkerHandler,
-    worker_pool::{SplitConfig, WorkerPoolHandler},
+    worker_pool::{WorkerPoolHandler, SplitConfig},
     Result,
 };
 

apps/sim-engine/lib/execution/src/runner/javascript/error.rs

@@ -91,6 +91,9 @@ pub enum JavaScriptError {
         None => String::new()
     })]
     JavascriptException(String, Option<String>),
+
+    #[error("Could not compile TypeScript file {filename}: {error}")]
+    TypeScriptCompilation { filename: String, error: String },
 }
 
 impl From<&str> for JavaScriptError {

apps/sim-engine/lib/execution/src/runner/javascript/modules.rs

@@ -42,6 +42,7 @@ pub(in crate::runner::javascript) fn import_module<'s>(
     let source_code = read_file(path).map_err(|err| {
         JavaScriptError::AccessJavascriptImport(path.to_string(), err.to_string())
     })?;
+
     let js_source_code = new_js_string(scope, &source_code);
     let js_path = new_js_string(scope, path);
     let source_map_url = v8::undefined(scope);

apps/sim-engine/lib/execution/src/runner/javascript/task.rs

@@ -13,7 +13,7 @@ pub(in crate::runner::javascript) fn get_next_task<'s>(
             let target = target.to_rust_string_lossy(scope);
 
             match target.as_str() {
-                "JavaScript" => MessageTarget::JavaScript,
+                "JavaScript" | "TypeScript" => MessageTarget::JavaScript,
                 "Python" => MessageTarget::Python,
                 "Rust" => MessageTarget::Rust,
                 "Dynamic" => MessageTarget::Dynamic,

apps/sim-engine/lib/execution/src/runner/language.rs

@@ -2,14 +2,15 @@ use core::fmt;
 
 use serde::{Deserialize, Serialize};
 
-use crate::{Error, Result};
+use crate::{runner::JavaScriptError, Error, Result};
 
 /// Supported languages
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub enum Language {
     JavaScript = 0,
     Python = 1,
     Rust = 2,
+    TypeScript = 3,
 }
 
 impl fmt::Display for Language {
@@ -19,9 +20,13 @@ impl fmt::Display for Language {
 }
 
 impl Language {
-    pub const NUM: usize = 3;
-    pub const ORDERED: [Language; Self::NUM] =
-        [Language::JavaScript, Language::Python, Language::Rust];
+    pub const NUM: usize = 4;
+    pub const ORDERED: [Language; Self::NUM] = [
+        Language::JavaScript,
+        Language::Python,
+        Language::Rust,
+        Language::TypeScript,
+    ];
 
     pub fn as_index(self) -> usize {
         self as usize
@@ -43,8 +48,76 @@ impl Language {
         match file_path.extension().and_then(std::ffi::OsStr::to_str) {
             Some("py") => Ok(Language::Python),
             Some("js") => Ok(Language::JavaScript),
+            Some("ts") => Ok(Language::TypeScript),
             Some("rs") => Ok(Language::Rust),
             _ => Err(Error::ParseBehavior(file_name.to_string())),
         }
     }
+
+    /// Add a compilation step if necessary, compiling from source code provided
+    /// by the user to the code that the runners can execute.
+    ///
+    /// TODO: Check validity of types
+    /// TODO: Add source map support
+    pub fn compile_source(&self, filename: &str, source: &str) -> Result<String> {
+        if *self != Language::TypeScript {
+            return Ok(source.to_string());
+        }
+
+        strip_typescript(filename, source)
+    }
+
+    /// Whether one language uses the same execution environment as another.
+    /// Useful for shared worker logic (the JavaScript worker is used for TypeScript support).
+    pub fn uses_same_runner(&self, other: &Language) -> bool {
+        match (self, other) {
+            (Language::JavaScript, Language::TypeScript) => true,
+            (Language::TypeScript, Language::JavaScript) => true,
+            _ => self == other,
+        }
+    }
+}
+
+fn strip_typescript(filename: &str, source_code: &str) -> Result<String> {
+    use std::sync::Arc;
+
+    use swc::config::Options;
+    use swc_common::{errors::Handler, Globals};
+
+    eprintln!(
+        "Stripping typescript from {filename} (source code: {})",
+        source_code.len()
+    );
+
+    let source_map = Arc::new(Default::default());
+    let compiler = swc::Compiler::new(Arc::clone(&source_map));
+    let source_file = source_map.new_source_file(
+        swc_common::FileName::Real(filename.into()),
+        source_code.into(),
+    );
+
+    let handler = Handler::with_emitter_writer(Box::new(vec![]), Some(source_map));
+
+    let options = Options::default();
+
+    swc_common::GLOBALS.set(&Globals::default(), || {
+        let s = compiler.process_js_file(source_file, &handler, &options);
+
+        match s {
+            Ok(v) => {
+                if handler.has_errors() {
+                    Err(Error::JavaScript(JavaScriptError::TypeScriptCompilation {
+                        filename: filename.to_string(),
+                        error: "Invalid TypeScript".to_string(),
+                    }))
+                } else {
+                    Ok(v.code)
+                }
+            }
+            Err(_e) => Err(Error::JavaScript(JavaScriptError::TypeScriptCompilation {
+                filename: filename.to_string(),
+                error: "Invalid TypeScript".to_string(),
+            })),
+        }
+    })
 }

apps/sim-engine/lib/execution/src/runner/target.rs

@@ -37,7 +37,7 @@ impl From<Language> for MessageTarget {
         match l {
             Language::Rust => Self::Rust,
             Language::Python => Self::Python,
-            Language::JavaScript => Self::JavaScript,
+            Language::JavaScript | Language::TypeScript => Self::JavaScript,
         }
     }
 }

apps/sim-engine/lib/experiment-structure/src/config/experiment.rs

@@ -37,7 +37,7 @@ impl ExperimentConfig {
         let package_config = PackageConfigBuilder::new()
             .add_init_package(match simulation.package_init.initial_state.name {
                 InitialStateName::InitJson => InitPackageName::Json,
-                InitialStateName::InitPy | InitialStateName::InitJs => InitPackageName::JsPy,
+                InitialStateName::InitPy | InitialStateName::InitJs | InitialStateName::InitTs=> InitPackageName::JsPy,
             })
             .build()?;
         let base_globals: Globals = serde_json::from_str(&simulation.globals_src)