From 5b61a75088065353433272b28015f6d141cc3bc0 Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Thu, 15 Aug 2024 16:44:42 -0400
Subject: [PATCH] feat: handler withdraw purpose

---
 crates/aiken-lang/src/ast.rs      | 22 ++++++++++++++++------
 crates/aiken-lang/src/builtins.rs |  4 ++++
 crates/aiken-lang/src/gen_uplc.rs | 23 ++++++++++++++++++++++-
 3 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/crates/aiken-lang/src/ast.rs b/crates/aiken-lang/src/ast.rs
index c65e9314..2cfd76ad 100644
--- a/crates/aiken-lang/src/ast.rs
+++ b/crates/aiken-lang/src/ast.rs
@@ -1,5 +1,8 @@
 use crate::{
-    builtins::{self, g1_element, g2_element},
+    builtins::{
+        self, g1_element, g2_element, SCRIPT_CONTEXT, SCRIPT_PURPOSE, SCRIPT_PURPOSE_MINT,
+        SCRIPT_PURPOSE_SPEND, SCRIPT_PURPOSE_WITHDRAW,
+    },
     expr::{TypedExpr, UntypedExpr},
     line_numbers::LineNumbers,
     parser::token::{Base, Token},
@@ -429,14 +432,21 @@ impl TypedDataType {
             constructors: vec![
                 RecordConstructor {
                     location: Span::empty(),
-                    name: "__Mint".to_string(),
+                    name: SCRIPT_PURPOSE_MINT.to_string(),
                     arguments: vec![],
                     doc: None,
                     sugar: false,
                 },
                 RecordConstructor {
                     location: Span::empty(),
-                    name: "__Spend".to_string(),
+                    name: SCRIPT_PURPOSE_SPEND.to_string(),
+                    arguments: vec![],
+                    doc: None,
+                    sugar: false,
+                },
+                RecordConstructor {
+                    location: Span::empty(),
+                    name: SCRIPT_PURPOSE_WITHDRAW.to_string(),
                     arguments: vec![],
                     doc: None,
                     sugar: false,
@@ -444,7 +454,7 @@ impl TypedDataType {
             ],
             doc: None,
             location: Span::empty(),
-            name: "__ScriptPurpose".to_string(),
+            name: SCRIPT_PURPOSE.to_string(),
             opaque: false,
             parameters: vec![],
             public: true,
@@ -456,14 +466,14 @@ impl TypedDataType {
         DataType {
             constructors: vec![RecordConstructor {
                 location: Span::empty(),
-                name: "__ScriptContext".to_string(),
+                name: SCRIPT_CONTEXT.to_string(),
                 arguments: vec![],
                 doc: None,
                 sugar: false,
             }],
             doc: None,
             location: Span::empty(),
-            name: "__ScriptContext".to_string(),
+            name: SCRIPT_CONTEXT.to_string(),
             opaque: false,
             parameters: vec![],
             public: true,
diff --git a/crates/aiken-lang/src/builtins.rs b/crates/aiken-lang/src/builtins.rs
index 3ee74864..72a6854c 100644
--- a/crates/aiken-lang/src/builtins.rs
+++ b/crates/aiken-lang/src/builtins.rs
@@ -40,6 +40,10 @@ pub const FUZZER: &str = "Fuzzer";
 pub const SCRIPT_PURPOSE: &str = "__ScriptPurpose";
 pub const SCRIPT_PURPOSE_SPEND: &str = "__Spend";
 pub const SCRIPT_PURPOSE_MINT: &str = "__Mint";
+pub const SCRIPT_PURPOSE_WITHDRAW: &str = "__Withdraw";
+pub const SCRIPT_PURPOSE_PUBLISH: &str = "__Publish";
+pub const SCRIPT_PURPOSE_VOTE: &str = "__Vote";
+pub const SCRIPT_PURPOSE_PROPOSE: &str = "__Propose";
 pub const SCRIPT_PURPOSES_COUNT: u16 = 6;
 
 pub const SCRIPT_CONTEXT: &str = "__ScriptContext";
diff --git a/crates/aiken-lang/src/gen_uplc.rs b/crates/aiken-lang/src/gen_uplc.rs
index 45f9da2e..40116b76 100644
--- a/crates/aiken-lang/src/gen_uplc.rs
+++ b/crates/aiken-lang/src/gen_uplc.rs
@@ -19,7 +19,7 @@ use crate::{
     },
     builtins::{
         bool, byte_array, data, function, int, list, option, pair, script_context, script_purpose,
-        void, PRELUDE, SCRIPT_PURPOSE_MINT, SCRIPT_PURPOSE_SPEND,
+        void, PRELUDE, SCRIPT_PURPOSE_MINT, SCRIPT_PURPOSE_SPEND, SCRIPT_PURPOSE_WITHDRAW,
     },
     expr::TypedExpr,
     gen_uplc::{
@@ -293,6 +293,27 @@ impl<'a> CodeGenerator<'a> {
                                         tipo: function(vec![byte_array()], script_purpose()),
                                     },
 
+                                    "withdraw" => TypedPattern::Constructor {
+                                        is_record: false,
+                                        location: Span::empty(),
+                                        name: SCRIPT_PURPOSE_WITHDRAW.to_string(),
+                                        arguments: vec![CallArg {
+                                            label: None,
+                                            location: Span::empty(),
+                                            value: TypedPattern::Var {
+                                                name: "__purpose_arg__".to_string(),
+                                                location: Span::empty(),
+                                            },
+                                        }],
+                                        module: None,
+                                        constructor: PatternConstructor::Record {
+                                            name: SCRIPT_PURPOSE_WITHDRAW.to_string(),
+                                            field_map: None,
+                                        },
+                                        spread_location: None,
+                                        tipo: function(vec![data()], script_purpose()),
+                                    },
+
                                     purpose => {
                                         unreachable!("unexpected/unknown purpose: {:?}", purpose)
                                     }