Add ascent_source! macro and include_source! syntax (#63)

as an `#include` style composability feature

Author: s-arash

README.MD

@@ -185,6 +185,29 @@ It may be useful to define macros that expand to either body items or head items
 
 You can find more about macros in Ascent macros [here](MACROS.MD).
 
+### `ascent_source!` and `include_source!`
+You can write Ascent code inside `ascent_source!` macro invocations and later include them in actual Ascent programs.
+This allows reusing your Ascent code across different Ascent programs, and composing Ascent programs from different code "modules".
+```Rust
+mod ascent_code {
+   ascent::ascent_source! { my_awesome_analysis:
+      // Ascent code for my awesome analysis
+   }
+}
+
+// elsewhere:
+ascent! {
+   struct MyAwesomeAnalysis;
+   include_source!(ascent_code::my_awesome_analysis);
+}
+
+ascent_par! {
+   struct MyAwesomeAnalysisParallelized;
+   include_source!(ascent_code::my_awesome_analysis);
+}
+
+```
+
 ### Misc
 - **`#![measure_rule_times]`** causes execution times of individual rules to be measured. Example: 
    ```Rust

ascent/examples/ascent_source.rs

@@ -0,0 +1,30 @@
+use ascent::{ascent_run, ascent_source};
+
+mod base {
+   ascent::ascent_source! {
+      /// Defines `edge` and `path`, the transitive closure of `edge`.
+      /// The type of a node is `Node`
+      tc:
+
+      relation edge(Node, Node);
+      relation path(Node, Node);
+      path(x, y) <-- edge(x, y);
+      path(x, z) <-- edge(x, y), path(y, z);
+   }
+}
+
+ascent_source! { symm_edges:
+   edge(x, y) <-- edge(y, x);
+}
+
+fn main() {
+   type Node = usize;
+   let res = ascent_run! {
+      include_source!(base::tc);
+      include_source!(symm_edges);
+
+      edge(1, 2), edge(2, 3), edge(3, 4);
+   };
+
+   assert!(res.path.contains(&(4, 1)));
+}

ascent/src/lib.rs

@@ -27,7 +27,7 @@ mod tuple_of_borrowed;
 mod rel_index_boilerplate;
 
 pub use ascent_base::*;
-pub use ascent_macro::{ascent, ascent_run};
+pub use ascent_macro::{ascent, ascent_run, ascent_source};
 #[cfg(feature = "par")]
 pub use ascent_macro::{ascent_par, ascent_run_par};
 #[cfg(feature = "par")]

ascent_macro/src/ascent_hir.rs

@@ -83,7 +83,6 @@ pub(crate) struct AscentIr {
    pub relations_ir_relations: HashMap<RelationIdentity, HashSet<IrRelation>>,
    pub relations_full_indices: HashMap<RelationIdentity, IrRelation>,
    pub lattices_full_indices: HashMap<RelationIdentity, IrRelation>,
-   // pub relations_no_indices: HashMap<RelationIdentity, IrRelation>,
    pub relations_metadata: HashMap<RelationIdentity, RelationMetadata>,
    pub rules: Vec<IrRule>,
    pub signatures: Signatures,
@@ -231,7 +230,6 @@ pub(crate) fn compile_ascent_program_to_hir(prog: &AscentProgram, is_parallel: b
    let mut relations_full_indices = HashMap::with_capacity(num_relations);
    let mut relations_initializations = HashMap::new();
    let mut relations_metadata = HashMap::with_capacity(num_relations);
-   // let mut relations_no_indices = HashMap::new();
    let mut lattices_full_indices = HashMap::new();
 
    let mut rel_identities = prog.relations.iter().map(|rel| (rel, RelationIdentity::from(rel))).collect_vec();
@@ -251,12 +249,11 @@ pub(crate) fn compile_ascent_program_to_hir(prog: &AscentProgram, is_parallel: b
       let rel_full_index = IrRelation::new(rel_identity.clone(), full_indices);
 
       relations_ir_relations.entry(rel_identity.clone()).or_default().insert(rel_full_index.clone());
-      // relations_ir_relations.entry(rel_identity.clone()).or_default().insert(rel_no_index.clone());
       relations_full_indices.insert(rel_identity.clone(), rel_full_index);
       if let Some(init_expr) = &rel.initialization {
          relations_initializations.insert(rel_identity.clone(), Rc::new(init_expr.clone()));
       }
-      
+
       let ds_attr = match (ds_attribute, rel.is_lattice) {
          (None, true) => None,
          (None, false) => Some(config.default_ds.clone()),
@@ -278,7 +275,6 @@ pub(crate) fn compile_ascent_program_to_hir(prog: &AscentProgram, is_parallel: b
          ),
          ds_attr,
       });
-      // relations_no_indices.insert(rel_identity, rel_no_index);
    }
    for (ir_rule, extra_relations) in ir_rules.iter() {
       for bitem in ir_rule.body_items.iter() {
@@ -303,7 +299,6 @@ pub(crate) fn compile_ascent_program_to_hir(prog: &AscentProgram, is_parallel: b
       relations_full_indices,
       lattices_full_indices,
       relations_metadata,
-      // relations_no_indices,
       signatures,
       config,
       is_parallel,
@@ -522,7 +517,7 @@ pub(crate) fn prog_get_relation<'a>(
                format!(
                   "wrong arity for relation `{name}` (expected {expected}, found {found})",
                   expected = rel.field_types.len(),
-                  found = arity, 
+                  found = arity,
                ),
             ))
          } else {

ascent_macro/src/ascent_mir.rs

@@ -312,7 +312,6 @@ pub(crate) fn compile_hir_to_mir(hir: &AscentIr) -> syn::Result<AscentMir> {
       relations_ir_relations: hir.relations_ir_relations.clone(),
       relations_full_indices: hir.relations_full_indices.clone(),
       lattices_full_indices: hir.lattices_full_indices.clone(),
-      // relations_no_indices: hir.relations_no_indices.clone(),
       relations_metadata: hir.relations_metadata.clone(),
       signatures: hir.signatures.clone(),
       config: hir.config.clone(),

ascent_macro/src/ascent_syntax.rs

@@ -5,17 +5,18 @@ use std::sync::Mutex;
 
 use ascent_base::util::update;
 use derive_syn_parse::Parse;
-use itertools::Itertools;
+use itertools::{Either, Itertools};
 use proc_macro2::{Span, TokenStream};
 use quote::ToTokens;
 use syn::parse::{Parse, ParseStream, Parser};
 use syn::punctuated::Punctuated;
 use syn::spanned::Spanned;
 use syn::{
-   Attribute, Error, Expr, ExprMacro, Generics, Ident, ImplGenerics, Pat, Result, Token, Type, TypeGenerics,
+   Attribute, Error, Expr, ExprMacro, Generics, Ident, ImplGenerics, Pat, Path, Result, Token, Type, TypeGenerics,
    Visibility, WhereClause, braced, parenthesized, parse2,
 };
 
+use crate::AscentMacroKind;
 use crate::syn_utils::{
    expr_get_vars, expr_visit_free_vars_mut, expr_visit_idents_in_macros_mut, pattern_get_vars, pattern_visit_vars_mut,
    token_stream_idents, token_stream_replace_ident,
@@ -45,13 +46,13 @@ mod kw {
    pub struct LongLeftArrow(Token![<], Token![-], Token![-]);
    #[allow(unused)]
    impl LongLeftArrow {
-      pub fn span(&self) -> Span {
-         join_spans([self.0.span, self.1.span, self.2.span])
-      }
+      pub fn span(&self) -> Span { join_spans([self.0.span, self.1.span, self.2.span]) }
    }
    syn::custom_keyword!(agg);
    syn::custom_keyword!(ident);
    syn::custom_keyword!(expr);
+
+   syn::custom_keyword!(include_source);
 }
 
 #[derive(Clone, Debug)]
@@ -117,7 +118,7 @@ fn parse_generics_with_where_clause(input: ParseStream) -> Result<Generics> {
    Ok(res)
 }
 
-// #[derive(Clone)]
+#[derive(PartialEq, Eq, Clone)]
 pub struct RelationNode {
    pub attrs: Vec<Attribute>,
    pub name: Ident,
@@ -126,6 +127,7 @@ pub struct RelationNode {
    pub _semi_colon: Token![;],
    pub is_lattice: bool,
 }
+
 impl Parse for RelationNode {
    fn parse(input: ParseStream) -> Result<Self> {
       let is_lattice = input.peek(kw::lattice);
@@ -549,6 +551,18 @@ pub struct MacroDefNode {
    body: TokenStream,
 }
 
+#[derive(Parse)]
+pub struct IncludeSourceNode {
+   pub include_source_kw: kw::include_source,
+   _bang: Token![!],
+   #[paren]
+   _arg_paren: syn::token::Paren,
+   #[inside(_arg_paren)]
+   #[call(syn::Path::parse_mod_style)]
+   path: syn::Path,
+   _semi: Token![;],
+}
+
 // #[derive(Clone)]
 pub(crate) struct AscentProgram {
    pub rules: Vec<RuleNode>,
@@ -558,40 +572,85 @@ pub(crate) struct AscentProgram {
    pub macros: Vec<MacroDefNode>,
 }
 
-impl Parse for AscentProgram {
-   fn parse(input: ParseStream) -> Result<Self> {
-      let attributes = Attribute::parse_inner(input)?;
-      let mut struct_attrs = Attribute::parse_outer(input)?;
-      let signatures = if input.peek(Token![pub]) || input.peek(Token![struct]) {
-         let mut signatures = Signatures::parse(input)?;
-         signatures.declaration.attrs = std::mem::take(&mut struct_attrs);
-         Some(signatures)
+/// The output that should be emitted when an `include_source!()` is encountered
+pub(crate) struct IncludeSourceMacroCall {
+   /// the encountered `include_source!()`
+   pub include_node: IncludeSourceNode,
+   pub before_tokens: TokenStream,
+   pub after_tokens: TokenStream,
+   pub ascent_macro_name: Path,
+}
+
+impl IncludeSourceMacroCall {
+   /// The output that should be emitted
+   pub fn macro_call_output(&self) -> TokenStream {
+      let Self { include_node, before_tokens, after_tokens, ascent_macro_name } = self;
+      let include_macro_callback = &include_node.path;
+      quote_spanned! {include_macro_callback.span()=>
+         #include_macro_callback! { {#ascent_macro_name}, {#before_tokens}, {#after_tokens} }
+      }
+   }
+}
+
+pub(crate) fn parse_ascent_program(
+   input: ParseStream, ascent_macro_name: Path,
+) -> Result<Either<AscentProgram, IncludeSourceMacroCall>> {
+   let input_clone = input.cursor();
+   let attributes = Attribute::parse_inner(input)?;
+   let mut struct_attrs = Attribute::parse_outer(input)?;
+   let signatures = if input.peek(Token![pub]) || input.peek(Token![struct]) {
+      let mut signatures = Signatures::parse(input)?;
+      signatures.declaration.attrs = std::mem::take(&mut struct_attrs);
+      Some(signatures)
+   } else {
+      None
+   };
+   let mut rules = vec![];
+   let mut relations = vec![];
+   let mut macros = vec![];
+   while !input.is_empty() {
+      let attrs =
+         if !struct_attrs.is_empty() { std::mem::take(&mut struct_attrs) } else { Attribute::parse_outer(input)? };
+      if input.peek(kw::relation) || input.peek(kw::lattice) {
+         let mut relation_node = RelationNode::parse(input)?;
+         relation_node.attrs = attrs;
+         relations.push(relation_node);
+      } else if input.peek(Token![macro]) {
+         if !attrs.is_empty() {
+            return Err(Error::new(attrs[0].span(), "unexpected attribute(s)"));
+         }
+         macros.push(MacroDefNode::parse(input)?);
+      } else if input.peek(kw::include_source) {
+         if !attrs.is_empty() {
+            return Err(Error::new(attrs[0].span(), "unexpected attribute(s)"));
+         }
+         let before_tokens = input_clone
+            .token_stream()
+            .into_iter()
+            .take_while(|tt| !spans_eq(&tt.span(), &input.span()))
+            .collect::<TokenStream>();
+         let include_node = IncludeSourceNode::parse(input)?;
+         let after_tokens: TokenStream = input.parse()?;
+         let include_source_macro_call =
+            IncludeSourceMacroCall { include_node, before_tokens, after_tokens, ascent_macro_name };
+         return Ok(Either::Right(include_source_macro_call));
       } else {
-         None
-      };
-      let mut rules = vec![];
-      let mut relations = vec![];
-      let mut macros = vec![];
-      while !input.is_empty() {
-         let attrs =
-            if !struct_attrs.is_empty() { std::mem::take(&mut struct_attrs) } else { Attribute::parse_outer(input)? };
-         if input.peek(kw::relation) || input.peek(kw::lattice) {
-            let mut relation_node = RelationNode::parse(input)?;
-            relation_node.attrs = attrs;
-            relations.push(relation_node);
-         } else if input.peek(Token![macro]) {
-            if !attrs.is_empty() {
-               return Err(Error::new(attrs[0].span(), "unexpected attribute(s)"));
-            }
-            macros.push(MacroDefNode::parse(input)?);
-         } else {
-            if !attrs.is_empty() {
-               return Err(Error::new(attrs[0].span(), "unexpected attribute(s)"));
-            }
-            rules.push(RuleNode::parse(input)?);
+         if !attrs.is_empty() {
+            return Err(Error::new(attrs[0].span(), "unexpected attribute(s)"));
          }
+         rules.push(RuleNode::parse(input)?);
+      }
+   }
+   Ok(Either::Left(AscentProgram { rules, relations, signatures, attributes, macros }))
+}
+
+impl Parse for AscentProgram {
+   fn parse(input: ParseStream) -> Result<Self> {
+      match parse_ascent_program(input, AscentMacroKind::default().macro_path(Span::call_site()))? {
+         Either::Left(parsed) => Ok(parsed),
+         Either::Right(include) =>
+            Err(Error::new(include.include_node.include_source_kw.span(), "Encountered `include_source!`")),
       }
-      Ok(AscentProgram { rules, relations, signatures, attributes, macros })
    }
 }