From 4452a9eefdacb4a2023756453d6098972f9d4e95 Mon Sep 17 00:00:00 2001
From: Anders Eie <strykejernene@gmail.com>
Date: Sat, 23 Apr 2022 10:37:08 +0200
Subject: [PATCH] Add support for visibility for fields

---
 src/field.rs     | 10 ++++++++++
 src/fields.rs    | 17 +++++++++++++++++
 src/function.rs  |  3 ++-
 src/impl.rs      |  1 +
 src/struct.rs    | 11 +++++++++++
 tests/codegen.rs | 21 +++++++++++++++++++++
 6 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/src/field.rs b/src/field.rs
index 769158f..b892890 100644
--- a/src/field.rs
+++ b/src/field.rs
@@ -14,6 +14,9 @@ pub struct Field {
 
     /// Field annotation
     pub annotation: Vec<String>,
+
+    /// The visibility of the field
+    pub visibility: Option<String>,
 }
 
 impl Field {
@@ -27,6 +30,7 @@ impl Field {
             ty: ty.into(),
             documentation: Vec::new(),
             annotation: Vec::new(),
+            visibility: None,
         }
     }
 
@@ -41,4 +45,10 @@ impl Field {
         self.annotation = annotation.iter().map(|ann| ann.to_string()).collect();
         self
     }
+
+    /// Set the visibility of the field
+    pub fn vis(&mut self, visibility: &str) -> &mut Self {
+        self.visibility = Some(visibility.to_string());
+        self
+    }
 }
diff --git a/src/fields.rs b/src/fields.rs
index 6878818..10cef0c 100644
--- a/src/fields.rs
+++ b/src/fields.rs
@@ -37,9 +37,23 @@ impl Fields {
             ty: ty.into(),
             documentation: Vec::new(),
             annotation: Vec::new(),
+            visibility: None,
         })
     }
 
+    pub fn new_named<T>(&mut self, name: &str, ty: T) -> &mut Field
+    where
+        T: Into<Type>,
+    {
+        self.named(name, ty);
+        if let Fields::Named(ref mut fields) = *self {
+            fields.last_mut().unwrap()
+        }
+        else {
+            unreachable!()
+        }
+    }
+
     pub fn tuple<T>(&mut self, ty: T) -> &mut Self
     where
         T: Into<Type>,
@@ -74,6 +88,9 @@ impl Fields {
                                 write!(fmt, "{}\n", ann)?;
                             }
                         }
+                        if let Some(visibility) = &f.visibility {
+                            write!(fmt, "{} ", visibility)?;
+                        }
                         write!(fmt, "{}: ", f.name)?;
                         f.ty.fmt(fmt)?;
                         write!(fmt, ",\n")?;
diff --git a/src/function.rs b/src/function.rs
index 827ec6c..afa7579 100644
--- a/src/function.rs
+++ b/src/function.rs
@@ -129,11 +129,12 @@ impl Function {
         self.args.push(Field {
             name: name.to_string(),
             ty: ty.into(),
-            // While a `Field` is used here, both `documentation`
+            // While a `Field` is used here, both `documentation`, `visibility`
             // and `annotation` does not make sense for function arguments.
             // Simply use empty strings.
             documentation: Vec::new(),
             annotation: Vec::new(),
+            visibility: None,
         });
 
         self
diff --git a/src/impl.rs b/src/impl.rs
index 6e411b5..91ff557 100644
--- a/src/impl.rs
+++ b/src/impl.rs
@@ -89,6 +89,7 @@ impl Impl {
             ty: ty.into(),
             documentation: Vec::new(),
             annotation: Vec::new(),
+            visibility: None,
         });
 
         self
diff --git a/src/struct.rs b/src/struct.rs
index 8feae8a..112c3f4 100644
--- a/src/struct.rs
+++ b/src/struct.rs
@@ -96,6 +96,17 @@ impl Struct {
         self
     }
 
+    /// Create a named field for the struct.
+    ///
+    /// A struct can either set named fields with this function or tuple fields
+    /// with `tuple_field`, but not both.
+    pub fn new_field<T>(&mut self, name: &str, ty: T) -> &mut Field
+    where
+        T: Into<Type>,
+    {
+        self.fields.new_named(name, ty)
+    }
+
     /// Add a tuple field to the struct.
     ///
     /// A struct can either set tuple fields with this function or named fields
diff --git a/tests/codegen.rs b/tests/codegen.rs
index 77e7346..9086fb8 100644
--- a/tests/codegen.rs
+++ b/tests/codegen.rs
@@ -603,5 +603,26 @@ enum IpAddrKind {
     V6,
 }"#;
 
+    assert_eq!(scope.to_string(), &expect[1..]);
+}
+
+#[test]
+fn struct_with_member_visibility() {
+    let mut scope = Scope::new();
+
+    let struct_description = scope.new_struct("Foo");
+
+    let mut bar = Field::new("bar", "usize");
+    bar.vis("pub");
+
+    struct_description.push_field(bar);
+    struct_description.new_field("baz", "i16").vis("pub(crate)");
+
+    let expect = r#"
+struct Foo {
+    pub bar: usize,
+    pub(crate) baz: i16,
+}"#;
+
     assert_eq!(scope.to_string(), &expect[1..]);
 }
\ No newline at end of file