Flatten fixes (#63)

Support for more flattening edge cases.

Resolves
foundry-rs/foundry#6936 (comment)
and
foundry-rs/foundry#6936 (comment)

We still don't rename stuff in yul blocks.

prb-math turned out to be a great codebase to test flattening against as
it has a lot of unusual patterns (qulified imports, user-defined
functions on user-defined types, etc). Currently, their contracts
flattened by this flattener are only having troubles with yul code
blocks

src/artifacts/ast/mod.rs

@@ -560,15 +560,12 @@ impl VariableDeclaration {
  }
 }
 
-/// Structured documentation (NatSpec).
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
-#[serde(untagged)]
-pub enum StructuredDocumentation {
- /// The documentation is provided in the form of an AST node.
- Parsed { text: String },
- /// The documentation is provided in the form of a string literal.
- Text(String),
-}
+ast_node!(
+ /// Structured documentation (NatSpec).
+ struct StructuredDocumentation {
+ text: String,
+ }
+);
 
 ast_node!(
  /// An override specifier.
@@ -1030,20 +1027,33 @@ ast_node!(
  /// A using for directive.
  struct UsingForDirective {
  #[serde(default, deserialize_with = "serde_helpers::default_for_null")]
- function_list: Vec<FunctionIdentifierPath>,
+ function_list: Vec<UsingForFunctionItem>,
  #[serde(default)]
  global: bool,
  library_name: Option<UserDefinedTypeNameOrIdentifierPath>,
  type_name: Option<TypeName>,
  }
 );
 
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum UsingForFunctionItem {
+ Function(FunctionIdentifierPath),
+ OverloadedOperator(OverloadedOperator),
+}
+
 /// A wrapper around [IdentifierPath] for the [UsingForDirective].
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct FunctionIdentifierPath {
  pub function: IdentifierPath,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct OverloadedOperator {
+ pub definition: IdentifierPath,
+ pub operator: String,
+}
+
 ast_node!(
  /// An import directive.
  struct ImportDirective {

src/artifacts/ast/visitor.rs

@@ -50,6 +50,8 @@ pub trait Visitor {
  fn visit_return(&mut self, _return: &Return) {}
  fn visit_inheritance_specifier(&mut self, _specifier: &InheritanceSpecifier) {}
  fn visit_modifier_invocation(&mut self, _invocation: &ModifierInvocation) {}
+ fn visit_inline_assembly(&mut self, _assembly: &InlineAssembly) {}
+ fn visit_external_assembly_reference(&mut self, _ref: &ExternalInlineAssemblyReference) {}
 }
 
 pub trait Walk {
@@ -244,10 +246,10 @@ impl_walk!(Statement, visit_statement, |statement, visitor| {
  Statement::Return(statement) => {
  statement.walk(visitor);
  }
- Statement::Break(_)
- | Statement::Continue(_)
- | Statement::InlineAssembly(_)
- | Statement::PlaceholderStatement(_) => {}
+ Statement::InlineAssembly(assembly) => {
+  assembly.walk(visitor);
+ }
+ Statement::Break(_) | Statement::Continue(_) | Statement::PlaceholderStatement(_) => {}
  }
 });
 
@@ -327,7 +329,7 @@ impl_walk!(UsingForDirective, visit_using_for, |directive, visitor| {
  library_name.walk(visitor);
  }
  for function in &directive.function_list {
- function.function.walk(visitor);
+ function.walk(visitor);
  }
 });
 
@@ -526,6 +528,14 @@ impl_walk!(ModifierInvocation, visit_modifier_invocation, |invocation, visitor|
  invocation.modifier_name.walk(visitor);
 });
 
+impl_walk!(InlineAssembly, visit_inline_assembly, |assembly, visitor| {
+ assembly.external_references.iter().for_each(|reference| {
+ reference.walk(visitor);
+ });
+});
+
+impl_walk!(ExternalInlineAssemblyReference, visit_external_assembly_reference);
+
 impl_walk!(ElementaryTypeName, visit_elementary_type_name);
 impl_walk!(Literal, visit_literal);
 impl_walk!(ImportDirective, visit_import_directive);
@@ -594,3 +604,18 @@ impl_walk!(ElementaryOrRawTypeName, |type_name, visitor| {
  ElementaryOrRawTypeName::Raw(_) => {}
  }
 });
+
+impl_walk!(UsingForFunctionItem, |item, visitor| {
+ match item {
+ UsingForFunctionItem::Function(func) => {
+ func.function.walk(visitor);
+ }
+ UsingForFunctionItem::OverloadedOperator(operator) => {
+ operator.walk(visitor);
+ }
+ }
+});
+
+impl_walk!(OverloadedOperator, |operator, visitor| {
+ operator.definition.walk(visitor);
+});