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);
+});