Extend shank instructions to work with Instruction struct variants (#32)

* extend shank instructions to have more than one field * parse multiple instruction fields * add test examples * chore: rename instructionArg(s) to arg(s) - they are inside an instruction thus that context is already provided - shorter names result in better generated code with solita Co-authored-by: Thorsten Lorenz <[email protected]>
metaplex-foundation · Aug 21, 2022 · a5bd471 · a5bd471
1 parent 08dc8e5
commit a5bd471
Show file tree

Hide file tree

Showing 10 changed files with 391 additions and 81 deletions.
diff --git a/shank-idl/src/idl_instruction.rs b/shank-idl/src/idl_instruction.rs
@@ -5,6 +5,7 @@ use heck::MixedCase;
 use serde::{Deserialize, Serialize};
 use shank_macro_impl::instruction::{
     Instruction, InstructionAccount, InstructionVariant,
+    InstructionVariantFields,
 };
 
 use crate::{idl_field::IdlField, idl_type::IdlType};
@@ -49,27 +50,48 @@ impl TryFrom<InstructionVariant> for IdlInstruction {
     fn try_from(variant: InstructionVariant) -> Result<Self> {
         let InstructionVariant {
             ident,
-            field_ty,
+            field_tys,
             accounts,
             discriminant,
         } = variant;
 
         let name = ident.to_string();
-        let args: Vec<IdlField> = if let Some(field_ty) = field_ty {
-            let name = if field_ty.kind.is_custom() {
-                field_ty.ident.to_string().to_mixed_case()
-            } else {
-                "instructionArgs".to_string()
-            };
-            let ty = IdlType::try_from(field_ty)?;
-            vec![IdlField {
-                name,
-                ty,
-                attrs: None,
-            }]
-        } else {
-            vec![]
+        let parsed_idl_fields: Result<Vec<IdlField>, Error> = match field_tys {
+            InstructionVariantFields::Named(args) => {
+                let mut parsed: Vec<IdlField> = vec![];
+                for (field_name, field_ty) in args.iter() {
+                    let ty = IdlType::try_from(field_ty.clone())?;
+                    parsed.push(IdlField {
+                        name: field_name.to_mixed_case(),
+                        ty,
+                        attrs: None,
+                    })
+                }
+                Ok(parsed)
+            }
+            InstructionVariantFields::Unnamed(args) => {
+                let mut parsed: Vec<IdlField> = vec![];
+                for (index, field_ty) in args.iter().enumerate() {
+                    let name = if args.len() == 1 {
+                        if field_ty.kind.is_custom() {
+                            field_ty.ident.to_string().to_mixed_case()
+                        } else {
+                            "args".to_string()
+                        }
+                    } else {
+                        format!("arg{}", index).to_string()
+                    };
+                    let ty = IdlType::try_from(field_ty.clone())?;
+                    parsed.push(IdlField {
+                        name,
+                        ty,
+                        attrs: None,
+                    })
+                }
+                Ok(parsed)
+            }
         };
+        let args: Vec<IdlField> = parsed_idl_fields?;
 
         let accounts = accounts.into_iter().map(IdlAccountItem::from).collect();
         ensure!(

diff --git a/shank-idl/tests/fixtures/instructions/single_file/instruction_with_args.json b/shank-idl/tests/fixtures/instructions/single_file/instruction_with_args.json
@@ -37,7 +37,7 @@
       ],
       "args": [
         {
-          "name": "instructionArgs",
+          "name": "args",
           "type": {
             "option": "u8"
           }

diff --git a/shank-idl/tests/fixtures/instructions/single_file/instruction_with_multiple_args.json b/shank-idl/tests/fixtures/instructions/single_file/instruction_with_multiple_args.json
@@ -0,0 +1,43 @@
+{
+  "version": "",
+  "name": "",
+  "instructions": [
+    {
+      "name": "CloseThing",
+      "accounts": [
+        {
+          "name": "creator",
+          "isMut": false,
+          "isSigner": true
+        }
+      ],
+      "args": [
+        {
+          "name": "arg0",
+          "type": {
+            "option": "u8"
+          }
+        },
+        {
+          "name": "arg1",
+          "type": {
+            "defined": "ComplexArgs"
+          }
+        },
+        {
+          "name": "arg2",
+          "type": {
+            "defined": "ComplexArgs"
+          }
+        }
+      ],
+      "discriminant": {
+        "type": "u8",
+        "value": 0
+      }
+    }
+  ],
+  "metadata": {
+    "origin": "shank"
+  }
+}
diff --git a/shank-idl/tests/fixtures/instructions/single_file/instruction_with_multiple_args.rs b/shank-idl/tests/fixtures/instructions/single_file/instruction_with_multiple_args.rs
@@ -0,0 +1,5 @@
+#[derive(ShankInstruction)]
+pub enum Instruction {
+    #[account(0, name = "creator", sig)]
+    CloseThing(Option<u8>, ComplexArgs, ComplexArgs),
+}
diff --git a/shank-idl/tests/fixtures/instructions/single_file/instruction_with_struct_args.json b/shank-idl/tests/fixtures/instructions/single_file/instruction_with_struct_args.json
@@ -0,0 +1,64 @@
+{
+  "version": "",
+  "name": "",
+  "instructions": [
+    {
+      "name": "CreateThing",
+      "accounts": [
+        {
+          "name": "creator",
+          "isMut": false,
+          "isSigner": true
+        },
+        {
+          "name": "thing",
+          "isMut": true,
+          "isSigner": false
+        }
+      ],
+      "args": [
+        {
+          "name": "someArgs",
+          "type": {
+            "defined": "SomeArgs"
+          }
+        },
+        {
+          "name": "otherArgs",
+          "type": {
+            "defined": "OtherArgs"
+          }
+        }
+      ],
+      "discriminant": {
+        "type": "u8",
+        "value": 0
+      }
+    },
+    {
+      "name": "CloseThing",
+      "accounts": [
+        {
+          "name": "creator",
+          "isMut": false,
+          "isSigner": true
+        }
+      ],
+      "args": [
+        {
+          "name": "args",
+          "type": {
+            "option": "u8"
+          }
+        }
+      ],
+      "discriminant": {
+        "type": "u8",
+        "value": 1
+      }
+    }
+  ],
+  "metadata": {
+    "origin": "shank"
+  }
+}
diff --git a/shank-idl/tests/fixtures/instructions/single_file/instruction_with_struct_args.rs b/shank-idl/tests/fixtures/instructions/single_file/instruction_with_struct_args.rs
@@ -0,0 +1,11 @@
+#[derive(ShankInstruction)]
+pub enum Instruction {
+    #[account(0, name = "creator", sig)]
+    #[account(1, name = "thing", mut)]
+    CreateThing {
+        some_args: SomeArgs,
+        other_args: OtherArgs,
+    },
+    #[account(0, name = "creator", sig)]
+    CloseThing(Option<u8>),
+}
diff --git a/shank-idl/tests/instructions.rs b/shank-idl/tests/instructions.rs
@@ -43,6 +43,40 @@ fn instruction_from_single_file_with_args() {
     assert_eq!(idl, expected_idl);
 }
 
+#[test]
+fn instruction_from_single_file_with_struct_args() {
+    let file = fixtures_dir()
+        .join("single_file")
+        .join("instruction_with_struct_args.rs");
+    let idl = parse_file(&file, &ParseIdlConfig::optional_program_address())
+        .expect("Parsing should not fail")
+        .expect("File contains IDL");
+
+    let expected_idl: Idl = serde_json::from_str(include_str!(
+        "./fixtures/instructions/single_file/instruction_with_struct_args.json"
+    ))
+    .unwrap();
+
+    assert_eq!(idl, expected_idl);
+}
+
+#[test]
+fn instruction_from_single_file_with_multiple_args() {
+    let file = fixtures_dir()
+        .join("single_file")
+        .join("instruction_with_multiple_args.rs");
+    let idl = parse_file(&file, &ParseIdlConfig::optional_program_address())
+        .expect("Parsing should not fail")
+        .expect("File contains IDL");
+
+    let expected_idl: Idl = serde_json::from_str(include_str!(
+        "./fixtures/instructions/single_file/instruction_with_multiple_args.json"
+    ))
+    .unwrap();
+
+    assert_eq!(idl, expected_idl);
+}
+
 #[test]
 fn instruction_from_single_file_with_optional_account() {
     let file = fixtures_dir()

diff --git a/shank-macro-impl/src/instruction/instruction.rs b/shank-macro-impl/src/instruction/instruction.rs
@@ -70,13 +70,19 @@ impl TryFrom<&ParsedEnum> for Instruction {
     }
 }
 
+#[derive(Debug)]
+pub enum InstructionVariantFields {
+    Unnamed(Vec<RustType>),
+    Named(Vec<(String, RustType)>),
+}
+
 // -----------------
 // Instruction Variant
 // -----------------
 #[derive(Debug)]
 pub struct InstructionVariant {
     pub ident: Ident,
-    pub field_ty: Option<RustType>,
+    pub field_tys: InstructionVariantFields,
     pub accounts: Vec<InstructionAccount>,
     pub discriminant: usize,
 }
@@ -93,19 +99,35 @@ impl TryFrom<&ParsedEnumVariant> for InstructionVariant {
             ..
         } = variant;
 
-        if fields.len() > 1 {
-            return Err(ParseError::new_spanned(
-                fields.get(1).map(|x| &x.rust_type.ident),
-                "An Instruction can only have one arg field",
-            ));
-        }
-        let field_ty = fields.first().map(|x| x.rust_type.clone());
+        let field_tys: InstructionVariantFields = if fields.len() > 0 {
+            // Determine if the InstructionType is tuple or struct variant
+            let field = fields.get(0).unwrap();
+            match &field.ident {
+                Some(_) => InstructionVariantFields::Named(
+                    fields
+                        .iter()
+                        .map(|x| {
+                            (
+                                x.ident.as_ref().unwrap().to_string(),
+                                x.rust_type.clone(),
+                            )
+                        })
+                        .collect(),
+                ),
+                None => InstructionVariantFields::Unnamed(
+                    fields.iter().map(|x| x.rust_type.clone()).collect(),
+                ),
+            }
+        } else {
+            InstructionVariantFields::Unnamed(vec![])
+        };
+
         let attrs: &[Attribute] = attrs.as_ref();
         let accounts: InstructionAccounts = attrs.try_into()?;
 
         Ok(Self {
             ident: ident.clone(),
-            field_ty,
+            field_tys,
             accounts: accounts.0,
             discriminant: *discriminant,
         })