helpers for reading ast Data and getting children of ast nodes

src/ast2/mod.rs

@@ -870,7 +870,7 @@ enum Tag {
     /// `data` is none
     ReturnStmt,
     /// `data` is an index to an expr
-    ReturnValueStmt,
+    ReturnExprStmt,
     /// `data` is a range from a..b into `extra` of an intern to a name and an optional intern to a type
     VarDecl,
     /// `data` is a range from a..b into `extra` of an intern to a name and an optional intern to a type
@@ -976,6 +976,9 @@ enum ParseError {
 pub struct Index(NonZero<u32>);
 
 impl Index {
+    pub fn new(i: u32) -> Index {
+        Self(NonZero::<u32>::new(i).unwrap())
+    }
     pub fn as_u32(&self) -> &u32 {
         unsafe { core::mem::transmute(self) }
     }
@@ -1010,6 +1013,42 @@ union Data {
     index_and_extra_offset: (Index, u32),
 }
 
+impl Data {
+    fn as_error(self) -> ParseError {
+        unsafe { self.error }
+    }
+    fn as_index(self) -> Index {
+        unsafe { self.index }
+    }
+    fn as_two_indices(self) -> (Index, Index) {
+        unsafe { self.two_indices }
+    }
+    fn as_index_range(self) -> (Index, Index) {
+        unsafe { self.range }
+    }
+    fn as_extra_range(self) -> (usize, usize) {
+        let (a, b) = unsafe { self.extra_range };
+        (a as usize, b as usize)
+    }
+    fn as_intern(self) -> Intern {
+        unsafe { self.intern }
+    }
+    fn as_two_indices(self) -> (Intern, Intern) {
+        unsafe { self.two_indices }
+    }
+    fn as_index_intern(self) -> (Index, intern::Index) {
+        unsafe { self.index_intern }
+    }
+    fn as_index_and_extra_offset(self) -> (Index, usize) {
+        let (i, e) = unsafe { self.index_and_extra_offset };
+        (i, e as usize)
+    }
+    fn as_intern_and_extra_offset(self) -> (intern::Index, usize) {
+        let (i, e) = unsafe { self.intern_and_extra_offset };
+        (i, e as usize)
+    }
+}
+
 impl Data {
     fn none() -> Self {
         Self { none: () }
@@ -1124,7 +1163,7 @@ impl Ast {
         let i = self.reserve_node();
         match expr {
             Some(expr) => {
-                self.set_tag_data_source_loc(i, Tag::ReturnValueStmt, Data::index(expr), loc)
+                self.set_tag_data_source_loc(i, Tag::ReturnExprStmt, Data::index(expr), loc)
             }
             None => self.set_tag_data_source_loc(i, Tag::ReturnStmt, Data::none(), loc),
         }
@@ -1145,7 +1184,8 @@ impl Ast {
         let start = self.extra.len() as u32;
         self.extra.push(name.into_u32());
         _ = self.extend_extra(assignment.map(|i| i.into_u32()));
-        let (_, end) = self.extend_extra(ty.map(|i| i.into_u32()));
+        _ = self.extend_extra(ty.map(|i| i.into_u32()));
+        let end = self.extra.len() as u32;
 
         let tag = match (is_let, assignment.is_some()) {
             (true, false) => Tag::VarDecl,
@@ -1387,6 +1427,111 @@ impl Ast {
     }
 }
 
+impl Ast {
+    fn get_node_children(&self, index: Index) -> Vec<Index> {
+        let tag = self.tags[index.index()];
+        let data = self.datas[index.index()];
+
+        match tag {
+            Tag::File => {
+                let (a, b) = data.as_extra_range();
+                self.extra[a..b].iter().map(|&i| Index::new(i)).collect()
+            }
+            Tag::FunctionProto => {
+                let (_, i) = data.as_intern_and_extra_offset();
+                vec![Index::new(self.extra[i + 1])]
+            }
+            Tag::FunctionDecl => {
+                let (a, b) = data.as_two_indices();
+                vec![a, b]
+            }
+            Tag::ParameterList => {
+                let (a, b) = data.as_extra_range();
+                self.extra[a..b].iter().map(|&i| Index::new(i)).collect()
+            }
+            Tag::Block | Tag::BlockTrailingExpr => {
+                let (a, b) = data.as_extra_range();
+                self.extra[a..b].iter().map(|&i| Index::new(i)).collect()
+            }
+            Tag::ExprStmt | Tag::ReturnExprStmt => {
+                let a = data.as_index();
+                vec![a]
+            }
+            Tag::Parameter => todo!(),
+            Tag::Constant => todo!(),
+            Tag::ReturnStmt => todo!(),
+            Tag::VarDecl => todo!(),
+            Tag::MutVarDecl => todo!(),
+            Tag::VarDeclAssignment | Tag::MutVarDeclAssignment => {
+                let (a, _) = data.as_extra_range();
+                let expr = Index::new(self.extra[a + 1]);
+                vec![expr]
+            }
+            Tag::GlobalDecl => {
+                let (a, _) = data.as_index_intern();
+                vec![a]
+            }
+            Tag::StructDecl => todo!(),
+            Tag::DeclRef => todo!(),
+            Tag::CallExpr => {
+                let (a, b) = data.as_two_indices();
+                vec![a, b]
+            }
+            Tag::ArgumentList => {
+                let (a, b) = data.as_extra_range();
+                self.extra[a..b].iter().map(|&i| Index::new(i)).collect()
+            }
+            Tag::Argument => {
+                let a = data.as_index();
+                vec![a]
+            }
+            Tag::NamedArgument => {
+                let (a, _) = data.as_index_intern();
+                vec![a]
+            }
+            Tag::ExplicitCast => {
+                let (a, _) = data.as_index_intern();
+                vec![a]
+            }
+            Tag::Deref | Tag::AddressOf | Tag::Not | Tag::Negate => {
+                let a = data.as_index();
+                vec![a]
+            }
+            Tag::Or
+            | Tag::And
+            | Tag::BitOr
+            | Tag::BitXOr
+            | Tag::BitAnd
+            | Tag::Eq
+            | Tag::NEq
+            | Tag::Lt
+            | Tag::Gt
+            | Tag::Le
+            | Tag::Ge
+            | Tag::Shl
+            | Tag::Shr
+            | Tag::Add
+            | Tag::Sub
+            | Tag::Mul
+            | Tag::Div
+            | Tag::Rem
+            | Tag::Assign
+            | Tag::SubscriptExpr
+            | Tag::IfExpr => {
+                let (a, b) = data.as_two_indices();
+                vec![a, b]
+            }
+            Tag::IfElseExpr => {
+                let (a, b) = data.as_index_and_extra_offset();
+                let if_ = Index::new(self.extra[b]);
+                let else_ = Index::new(self.extra[b + 1]);
+                vec![a, if_, else_]
+            }
+            _ => vec![],
+        }
+    }
+}
+
 mod ast_gen {
 
     use intern::{PointerFlags, SimpleType};
@@ -1394,7 +1539,7 @@ mod ast_gen {
     use num_bigint::{BigInt, BigUint};
 
     use crate::{
-        common::{from_lo_hi_dwords, NextIf},
+        common::from_lo_hi_dwords,
         comptime,
         lexer::{Radix, TokenItem, TokenIterator},
         tokens::PRECEDENCE_MAP,
@@ -2449,17 +2594,6 @@ mod ast_gen {
             }
         }
 
-        fn parse_statement(&mut self, tokens: &mut TokenIterator) -> ParseResult<Index> {
-            let expr = self.parse_expr(tokens)?;
-
-            let stmt = match tokens.eat_token(Token::Semi) {
-                Some(_) => self.ast.push_expr_stmt(expr),
-                None => self.push_error(ParseError::ExpectedStatement, self.ast.get_loc(expr)),
-            };
-
-            Ok(stmt)
-        }
-
         /// SUMTYPE_DECL <-
         ///     type IDENTIFIER = TYPE_UNION
         /// TYPE_UNION <-