Merge pull request #632 from powdr-labs/std

Add Powdr standard library

.github/workflows/pr-tests.yml

@@ -46,7 +46,7 @@ jobs:
     - name: Run default tests
       run: PILCOM=$(pwd)/pilcom/ cargo test --all --all-features --profile pr-tests --verbose
     - name: Run slow tests
-      # Number threads is set to 1 because the runner does not have enough memeory for more.
+      # Number threads is set to 1 because the runner does not have enough memory for more.
       run: PILCOM=$(pwd)/pilcom/ cargo test --all --all-features --profile pr-tests --verbose -- --ignored --nocapture --test-threads=1 --exact test_keccak test_vec_median instruction_tests::addi
     - name: Check benches compile without running them
       run: cargo bench --all --all-features --profile pr-tests --no-run

airgen/src/lib.rs

@@ -21,10 +21,16 @@ use number::FieldElement;
 pub fn compile<T: FieldElement>(input: AnalysisASMFile<T>) -> PILGraph<T> {
     let main_location = Location::main();
 
+    let non_std_machines = input
+        .machines
+        .iter()
+        .filter(|(k, _)| k.parts[0] != "std")
+        .collect::<BTreeMap<_, _>>();
+
     // we start from the main machine
-    let main_ty = match input.machines.len() {
+    let main_ty = match non_std_machines.len() {
         // if there is a single machine, treat it as main
-        1 => input.machines.keys().next().unwrap().clone(),
+        1 => (*non_std_machines.keys().next().unwrap()).clone(),
         // otherwise, use the machine called `MAIN`
         _ => {
             assert!(input

book/src/asm/modules.md

@@ -4,4 +4,6 @@ powdr exposes a module system to help organise and reuse code.
 
 ```
 {{#include ../../../test_data/asm/book/modules.asm}}
-```
+```
+
+Note that a module can't be called `std`, as this name is reserved for an upcoming powdr standard library.

importer/src/lib.rs

@@ -1,5 +1,6 @@
 mod module_loader;
 mod path_canonicalizer;
+mod powdr_std;
 
 use std::path::PathBuf;
 
@@ -8,12 +9,15 @@ pub use module_loader::load_module_files;
 use number::FieldElement;
 use parser::parse_asm;
 use path_canonicalizer::canonicalize_paths;
+use powdr_std::add_std;
 
 pub fn resolve<T: FieldElement>(
     path: Option<PathBuf>,
     module: ASMProgram<T>,
 ) -> Result<ASMProgram<T>, String> {
-    load_module_files(path, module).and_then(canonicalize_paths)
+    load_module_files(path, module)
+        .and_then(add_std)
+        .and_then(canonicalize_paths)
 }
 
 /// A test utility to process a source file until after import resolution

importer/src/powdr_std.rs

@@ -0,0 +1,115 @@
+use std::{env, path::PathBuf};
+
+use ast::parsed::{
+    asm::{
+        ASMModule, ASMProgram, Import, Module, ModuleStatement, Part, SymbolDefinition, SymbolPath,
+        SymbolValue,
+    },
+    folder::Folder,
+};
+use number::FieldElement;
+use parser::parse_asm;
+
+use crate::load_module_files;
+
+static POWDR_STD_ENV: &str = "POWDR_STD";
+static MOD_FILE: &str = "mod.asm";
+
+/// Loads the standard library module from the location specified in the <POWDR_STD_ENV> environment variable
+/// (or, if unset, <project_root>/std).
+///
+/// # Panics
+/// If there is an error loading the standard library
+fn load_std<T: FieldElement>() -> ASMModule<T> {
+    let default_std_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .parent()
+        .unwrap()
+        .join("std");
+    let std_path = env::var(POWDR_STD_ENV)
+        .map(PathBuf::from)
+        .unwrap_or(default_std_path);
+    let std_path = std_path.join(MOD_FILE);
+
+    match std::fs::read_to_string(&std_path) {
+        Err(_) => {
+            panic!(
+                "Couldn't find the powdr standard library at {}. Set the {} environment variable.",
+                std_path.display(),
+                POWDR_STD_ENV
+            )
+        }
+        Ok(std_source) => {
+            let std_content =
+                parse_asm::<T>(Some(std_path.as_path().to_str().unwrap()), &std_source).unwrap();
+
+            // This resolves all submodules and returns the standard library's main module
+            load_module_files(Some(std_path), std_content).unwrap().main
+        }
+    }
+}
+
+pub fn add_std<T: FieldElement>(program: ASMProgram<T>) -> Result<ASMProgram<T>, String> {
+    StdAdder().fold_program(program)
+}
+
+struct StdAdder();
+
+type Error = String;
+
+impl<T: FieldElement> Folder<T> for StdAdder {
+    type Error = Error;
+
+    fn fold_program(&mut self, p: ASMProgram<T>) -> Result<ASMProgram<T>, Self::Error> {
+        // Add `std` to the main module
+        let mut main = p.main;
+        main.statements
+            .push(ModuleStatement::SymbolDefinition(SymbolDefinition {
+                name: "std".to_string(),
+                value: SymbolValue::Module(Module::Local(load_std())),
+            }));
+
+        // Recurse
+        let main = self.fold_module_value(main)?;
+        Ok(ASMProgram { main })
+    }
+
+    fn fold_module_value(&mut self, module: ASMModule<T>) -> Result<ASMModule<T>, Self::Error> {
+        // This block is identical to Folder::fold_module_value.
+        // Unfortunately, there is no way to call the super method from here.
+        let mut statements = module
+            .statements
+            .into_iter()
+            .map(|s| match s {
+                ModuleStatement::SymbolDefinition(d) => match d.value {
+                    SymbolValue::Machine(machine) => self.fold_machine(machine).map(From::from),
+                    SymbolValue::Import(import) => {
+                        <StdAdder as Folder<T>>::fold_import(self, import).map(From::from)
+                    }
+                    SymbolValue::Module(module) => self.fold_module(module).map(From::from),
+                }
+                .map(|value| ModuleStatement::SymbolDefinition(SymbolDefinition { value, ..d })),
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+
+        // Check whether the module already has a definition for `std`
+        // (E.g. the main module)
+        let has_std = statements.iter().any(|s| match s {
+            ModuleStatement::SymbolDefinition(d) => d.name == "std",
+        });
+
+        if !has_std {
+            // If not, add `use super::std;`
+            let std_import_path = SymbolPath {
+                parts: [Part::Super, Part::Named("std".to_string())].into(),
+            };
+            statements.push(ModuleStatement::SymbolDefinition(SymbolDefinition {
+                name: "std".to_string(),
+                value: SymbolValue::Import(Import {
+                    path: std_import_path,
+                }),
+            }));
+        }
+
+        Ok(ASMModule { statements })
+    }
+}

std/hash/mod.asm

@@ -0,0 +1 @@
+mod poseidon_bn254;

std/mod.asm

@@ -0,0 +1 @@
+mod hash;

test_data/asm/poseidon_bn254_test.asm

@@ -1,6 +1,4 @@
-mod poseidon_bn254;
-
-use poseidon_bn254::PoseidonBN254;
+use std::hash::poseidon_bn254::PoseidonBN254;
 
 machine Main {
     degree 512;