added try statement as primitive

Cyfrin · Oct 6, 2024 · a343b3b · a343b3b
1 parent d05680e
commit a343b3b
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 1 deletion.
diff --git a/aderyn_core/src/context/flow/display.rs b/aderyn_core/src/context/flow/display.rs
@@ -19,6 +19,7 @@ impl CfgNodeDescriptor {
             CfgNodeDescriptor::EmitStatement(n) => n.peek(context),
             CfgNodeDescriptor::RevertStatement(n) => n.peek(context),
             CfgNodeDescriptor::InlineAssembly(n) => n.peek(context),
+            CfgNodeDescriptor::TryStatement(n) => n.peek(context),
             CfgNodeDescriptor::IfStatementCondition(n) => n.peek(context),
             CfgNodeDescriptor::WhileStatementCondition(n) => n.peek(context),
             CfgNodeDescriptor::ForStatementCondition(n) => n.peek(context),
@@ -246,3 +247,15 @@ impl CfgDoWhileStatementCondition {
         content
     }
 }
+
+impl CfgTryStatement {
+    pub fn peek(&self, context: &WorkspaceContext) -> String {
+        let mut content = format!("Try Stmt ({})", self.try_statement);
+        if let Some(node) = context.nodes.get(&self.try_statement) {
+            if let Some(inside) = node.peek(context) {
+                content.push_str(&format!(": \n{}", inside));
+            }
+        }
+        content
+    }
+}
diff --git a/aderyn_core/src/context/flow/kind.rs b/aderyn_core/src/context/flow/kind.rs
@@ -64,6 +64,7 @@ impl CfgNodeDescriptor {
             CfgNodeDescriptor::EmitStatement(_) => CfgNodeKind::Primitive,
             CfgNodeDescriptor::RevertStatement(_) => CfgNodeKind::Primitive,
             CfgNodeDescriptor::InlineAssembly(_) => CfgNodeKind::Primitive,
+            CfgNodeDescriptor::TryStatement(_) => CfgNodeKind::Primitive,
             CfgNodeDescriptor::IfStatementCondition(_) => CfgNodeKind::Primitive,
             CfgNodeDescriptor::WhileStatementCondition(_) => CfgNodeKind::Primitive,
             CfgNodeDescriptor::ForStatementCondition(_) => CfgNodeKind::Primitive,

diff --git a/aderyn_core/src/context/flow/mod.rs b/aderyn_core/src/context/flow/mod.rs
@@ -63,6 +63,7 @@ pub enum CfgNodeDescriptor {
     Return(Box<CfgReturnStatement>),
     EmitStatement(Box<CfgEmitStatement>),
     RevertStatement(Box<CfgRevertStatement>),
+    TryStatement(Box<CfgTryStatement>),
     InlineAssembly(Box<CfgInlineAssemblyStatement>),
     IfStatementCondition(Box<CfgIfStatementCondition>),
     WhileStatementCondition(Box<CfgWhileStatementCondition>),
@@ -260,6 +261,7 @@ impl Cfg {
             | CfgNodeDescriptor::RevertStatement(_)
             | CfgNodeDescriptor::PlaceholderStatement(_)
             | CfgNodeDescriptor::InlineAssembly(_)
+            | CfgNodeDescriptor::TryStatement(_)
             | CfgNodeDescriptor::EmitStatement(_)
             | CfgNodeDescriptor::ExpressionStatement(_)
             | CfgNodeDescriptor::WhileStatementCondition(_)

diff --git a/aderyn_core/src/context/flow/primitives.rs b/aderyn_core/src/context/flow/primitives.rs
@@ -264,3 +264,22 @@ impl Cfg {
         )))
     }
 }
+
+///////////////////////////////////////////////////////////////////////////
+
+#[derive(Debug, Clone)]
+pub struct CfgTryStatement {
+    pub try_statement: AstNodeId,
+}
+
+impl CfgTryStatement {
+    pub fn from(try_stmt: &TryStatement) -> Self {
+        Self { try_statement: try_stmt.id }
+    }
+}
+
+impl Cfg {
+    pub fn add_try_statement(&mut self, try_stmt: &TryStatement) -> CfgNodeId {
+        self.add_node(CfgNodeDescriptor::TryStatement(Box::new(CfgTryStatement::from(try_stmt))))
+    }
+}
diff --git a/aderyn_core/src/context/flow/utils.rs b/aderyn_core/src/context/flow/utils.rs
@@ -5,7 +5,7 @@ use super::{voids::CfgEndNode, AstNodeId, Cfg, CfgNodeDescriptor, CfgNodeId, Sta
 impl Cfg {
     pub fn add_statement_node(&mut self, stmt: &Statement) -> CfgNodeId {
         match stmt {
-            Statement::TryStatement(_) => unimplemented!(),
+            Statement::TryStatement(n) => self.add_try_statement(n),
             Statement::UncheckedBlock(n) => self.add_unchecked_block_node(n),
             Statement::DoWhileStatement(n) => self.add_do_while_statement(n),
             Statement::ForStatement(n) => self.add_for_statement(n),
@@ -209,6 +209,7 @@ impl CfgNodeDescriptor {
             CfgNodeDescriptor::WhileStatement(n) => Some(n.while_statement),
             CfgNodeDescriptor::ForStatement(n) => Some(n.for_statement),
             CfgNodeDescriptor::DoWhileStatement(n) => Some(n.do_while_statement),
+            CfgNodeDescriptor::TryStatement(n) => Some(n.try_statement),
         }
     }
 }