1 files changed, 42 insertions, 5 deletions

diff --git a/src/parse.rs b/src/parse.rs
index cbc37bf..b1721be 100644
--- a/src/parse.rs
+++ b/src/parse.rs
@@ -27,22 +27,24 @@ pub trait Expander {
 
 #[derive(Debug, Clone)]
 pub enum Ast<T: Stage> {
-    VarAssign(VarAssign),
+    FunDecl(FunDecl<T>),
+    VarAssign(VarAssign<T>),
     Pipes(Pipes<T>),
 }
 
 impl Ast<PreExpansion> {
     pub fn expand<E: Expander>(self, e: &mut E) -> Res<Ast<PostExpansion>, E::Error> {
         match self {
-            Ast::VarAssign(_) => todo!(),
+            Ast::VarAssign(va) => Ok(Ast::VarAssign(va.expand(e)?)),
             Ast::Pipes(pipes) => Ok(Ast::Pipes(pipes.expand(e)?)),
+            Ast::FunDecl(fd) => Ok(Ast::FunDecl(fd.expand(e)?)),
         }
     }
 }
 
 #[derive(Debug, Clone)]
 pub struct FunBody {
-    pub body: Ast<PreExpansion>,
+    pub body: Box<Ast<PreExpansion>>,
 }
 
 impl Parse for FunBody {
@@ -58,7 +60,7 @@ impl Parse for FunBody {
         }
 
         b.adv();
-        let body = Ast::parse(b)?;
+        let body = Box::new(Ast::parse(b)?);
         if b.is_empty() {
             if b.is_completion() {
                 Ok(Self { body })
@@ -81,7 +83,6 @@ pub struct FunDecl<S: Stage> {
 
 impl Parse for FunDecl<PreExpansion> {
     fn parse(b: &mut Cursor<'_>) -> Result<Self> {
-        b.spaces();
         if !b.buf.starts_with(b"fun ") && !b.buf.starts_with(b"fun\t") {
             return Err(ParseError::NotAFunDecl);
         }
@@ -93,6 +94,15 @@ impl Parse for FunDecl<PreExpansion> {
     }
 }
 
+impl FunDecl<PreExpansion> {
+    fn expand<E: Expander>(self, e: &mut E) -> Res<FunDecl<PostExpansion>, E::Error> {
+        Ok(FunDecl {
+            name: self.name.expand(e)?,
+            body: self.body,
+        })
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct VarDecl<S: Stage> {
     pub name: S::Str,
@@ -128,6 +138,15 @@ impl Parse for VarAssign<PreExpansion> {
     }
 }
 
+impl VarAssign<PreExpansion> {
+    fn expand<E: Expander>(self, e: &mut E) -> Res<VarAssign<PostExpansion>, E::Error> {
+        Ok(VarAssign {
+            var: self.var.expand(e)?,
+            val: self.val.expand(e)?,
+        })
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct Pipes<T: Stage> {
     pub cmds: Vec<Command<T>>,
@@ -550,6 +569,24 @@ impl<'a> Cursor<'a> {
 
 impl Parse for Ast<PreExpansion> {
     fn parse(b: &mut Cursor<'_>) -> Result<Self> {
+        b.spaces();
+
+        let orig_len = b.buf.len();
+        let x = VarAssign::parse(b);
+        if let Ok(va) = x {
+            return Ok(Self::VarAssign(va));
+        } else if b.buf.len() != orig_len {
+            x?;
+        }
+
+        let orig_len = b.buf.len();
+        let x = FunDecl::parse(b);
+        if let Ok(fd) = x {
+            return Ok(Self::FunDecl(fd));
+        } else if b.buf.len() != orig_len {
+            x?;
+        }
+
         Ok(Self::Pipes(b.parse()?))
     }
 }