yul func support

crates/graph/src/graph_elements.rs

@@ -103,6 +103,8 @@ pub enum Node {
     Msg(Msg),
     /// The `block` global in solidity
     Block(Block),
+    /// A yul-based function
+    YulFunction(YulFunction),
 }
 
 pub fn as_dot_str(
@@ -276,6 +278,8 @@ pub enum Edge {
     BuiltinFunction,
     /// A connection from one contract to another contract
     UsingContract(Loc),
+    /// Yul function
+    YulFunction(usize),
 }
 
 impl Heirarchical for Edge {
@@ -286,8 +290,8 @@ impl Heirarchical for Edge {
             Part | Import => 1,
 
             Contract | Ty | Field | Enum | Struct | Error | Event | Var | InheritedContract
-            | Modifier | FallbackFunc | Constructor | ReceiveFunc | LibraryFunction(_)
-            | BuiltinFunction | Func | UsingContract(_) => 2,
+            | Modifier | FallbackFunc | Constructor | YulFunction(_) | ReceiveFunc
+            | LibraryFunction(_) | BuiltinFunction | Func | UsingContract(_) => 2,
 
             Context(_) | ErrorParam | FunctionParam | FunctionReturn | FuncModifier(_) => 3,
         }

crates/graph/src/nodes/func_ty.rs

@@ -1,7 +1,7 @@
 use crate::{
     nodes::{
         Concrete, ContextNode, ContextVar, ContextVarNode, ContractNode, SourceUnitNode,
-        SourceUnitPartNode,
+        SourceUnitPartNode, YulFunctionNode,
     },
     range::elem::Elem,
     AnalyzerBackend, AsDotStr, ContextEdge, Edge, GraphBackend, Node, SolcRange, VarType,
@@ -53,6 +53,23 @@ impl FunctionNode {
         Ok(())
     }
 
+    pub fn yul_funcs(
+        &self,
+        analyzer: &impl GraphBackend,
+        assembly_block_idx: usize,
+    ) -> Vec<YulFunctionNode> {
+        analyzer
+            .graph()
+            .edges_directed(self.0.into(), Direction::Incoming)
+            .filter_map(|edge| match edge.weight() {
+                Edge::YulFunction(asm_block) if *asm_block == assembly_block_idx => {
+                    Some(edge.source().into())
+                }
+                _ => None,
+            })
+            .collect::<Vec<_>>()
+    }
+
     pub fn ty(&self, analyzer: &impl GraphBackend) -> Result<FunctionTy, GraphError> {
         Ok(self.underlying(analyzer)?.ty)
     }

crates/graph/src/nodes/mod.rs

@@ -42,3 +42,6 @@ pub use source_unit::*;
 
 mod debug_reconstruction;
 pub use debug_reconstruction::*;
+
+mod yul_func;
+pub use yul_func::*;

crates/graph/src/nodes/yul_func.rs

@@ -0,0 +1,79 @@
+use crate::{
+    nodes::{Concrete, ContractNode},
+    range::elem::Elem,
+    AnalyzerBackend, AsDotStr, Edge, GraphBackend, Node, VarType,
+};
+
+use shared::{FlatExpr, GraphError, NodeIdx, RangeArena};
+
+use petgraph::visit::EdgeRef;
+use solang_parser::pt::{Expression, Identifier, Loc, TypeDefinition};
+
+#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
+pub struct YulFunctionNode(pub usize);
+impl YulFunctionNode {
+    pub fn underlying<'a>(
+        &self,
+        analyzer: &'a impl GraphBackend,
+    ) -> Result<&'a YulFunction, GraphError> {
+        match analyzer.node(*self) {
+            Node::YulFunction(ty) => Ok(ty),
+            Node::Unresolved(ident) => Err(GraphError::UnknownVariable(format!(
+                "Could not find variable: {}",
+                ident.name
+            ))),
+            e => Err(GraphError::NodeConfusion(format!(
+                "Node type confusion: expected node to be YulFunctionNode but it was: {e:?}"
+            ))),
+        }
+    }
+
+    pub fn name(&self, analyzer: &impl GraphBackend) -> Result<String, GraphError> {
+        Ok(self.underlying(analyzer)?.name.to_string())
+    }
+    pub fn exprs(&self, analyzer: &impl GraphBackend) -> Result<Vec<FlatExpr>, GraphError> {
+        Ok(self.underlying(analyzer)?.exprs.clone())
+    }
+}
+
+impl From<YulFunctionNode> for NodeIdx {
+    fn from(val: YulFunctionNode) -> Self {
+        val.0.into()
+    }
+}
+
+impl From<NodeIdx> for YulFunctionNode {
+    fn from(idx: NodeIdx) -> Self {
+        YulFunctionNode(idx.index())
+    }
+}
+
+impl AsDotStr for YulFunctionNode {
+    fn as_dot_str(
+        &self,
+        analyzer: &impl GraphBackend,
+        _arena: &mut RangeArena<Elem<Concrete>>,
+    ) -> String {
+        let underlying = self.underlying(analyzer).unwrap();
+        format!("yul function {}", underlying.name,)
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct YulFunction {
+    pub loc: Loc,
+    pub name: &'static str,
+    pub exprs: Vec<FlatExpr>,
+}
+
+impl From<YulFunction> for Node {
+    fn from(val: YulFunction) -> Self {
+        Node::YulFunction(val)
+    }
+}
+
+impl YulFunction {
+    pub fn new(exprs: Vec<FlatExpr>, name: &'static str, loc: Loc) -> YulFunction {
+        YulFunction { loc, name, exprs }
+    }
+}

crates/graph/src/var_type.rs

@@ -248,7 +248,8 @@ impl VarType {
             | Node::Entry
             | Node::Context(..)
             | Node::Msg(_)
-            | Node::Block(_) => None,
+            | Node::Block(_)
+            | Node::YulFunction(..) => None,
         }
     }
 

crates/pyrometer/src/analyzer.rs

@@ -149,6 +149,7 @@ pub struct Analyzer {
 
     pub flattened: Vec<FlatExpr>,
     pub expr_flag: Option<ExprFlag>,
+    pub current_asm_block: usize,
 }
 
 impl<'a> Default for Analyzer {
@@ -189,6 +190,7 @@ impl<'a> Default for Analyzer {
             minimize_debug: None,
             flattened: vec![],
             expr_flag: None,
+            current_asm_block: 0,
         };
         a.builtin_fn_inputs = builtin_fns::builtin_fns_inputs(&mut a);
 

crates/pyrometer/src/analyzer_backend.rs

@@ -564,6 +564,17 @@ impl AnalyzerLike for Analyzer {
         self.flattened.push(flat);
     }
 
+    fn decrement_asm_block(&mut self) {
+        self.current_asm_block -= 1;
+    }
+    fn increment_asm_block(&mut self) {
+        self.current_asm_block += 1;
+    }
+
+    fn current_asm_block(&self) -> usize {
+        self.current_asm_block
+    }
+
     fn expr_stack(&self) -> &[FlatExpr] {
         &self.flattened
     }

crates/shared/src/analyzer_like.rs

@@ -197,6 +197,9 @@ pub trait AnalyzerLike: GraphLike {
     type FlatExpr;
     type ExprFlag: Copy;
     fn push_expr(&mut self, flat: Self::FlatExpr);
+    fn increment_asm_block(&mut self);
+    fn decrement_asm_block(&mut self);
+    fn current_asm_block(&self) -> usize;
     fn expr_stack(&self) -> &[Self::FlatExpr];
     fn expr_stack_mut(&mut self) -> &mut Vec<Self::FlatExpr>;
     fn take_expr_flag(&mut self) -> Option<Self::ExprFlag>;