New p2wpkh methods

src/bitcoin/Address.mo

@@ -1,15 +1,41 @@
 import P2pkh "./P2pkh";
+import P2WPKH "./P2wpkh";
+import P2WSH "./P2wsh";
 import P2tr "./P2tr";
 import Script "./Script";
 import Segwit "../Segwit";
 import Types "./Types";
 import Result "mo:base/Result";
+import Nat8 "mo:base/Nat8";
 
 module {
   public func addressFromText(address : Text) : Result.Result<Types.Address, Text> {
     switch (Segwit.decode(address)) {
-      case (#ok _) {
-        return #ok(#p2tr_key(address));
+      case (#ok((_hrp, witnessProgram))) {
+        // Successfully decoded a Bech32/Bech32m address, now check version/size
+        if (witnessProgram.version == 0) {
+          if (witnessProgram.program.size() == 20) {
+            // Version 0, 20-byte program -> P2WPKH
+            return #ok(#p2wpkh(address));
+          } else if (witnessProgram.program.size() == 32) {
+            // Version 0, 32-byte program -> P2WSH (Add #p2wsh variant to Types.mo first)
+            return #ok(#p2wsh(address)); // Enable if #p2wsh is added
+          } else {
+            return #err("Invalid program size for witness version 0");
+          };
+        } else if (witnessProgram.version == 1) {
+          if (witnessProgram.program.size() == 32) {
+            // Version 1, 32-byte program -> P2TR (Taproot)
+            // For now, assume key path spend as default when parsing address text
+            return #ok(#p2tr_key(address));
+            // TODO: Decide how to differentiate between P2TR key/script from address text alone if needed
+          } else {
+            return #err("Invalid program size for witness version 1");
+          };
+        } else {
+          // Other witness versions (>= 2) are currently unassigned by BIPs
+          return #err("Unsupported witness version: " # Nat8.toText(witnessProgram.version));
+        };
       };
       case (_) {};
     };
@@ -28,9 +54,15 @@ module {
   public func scriptPubKey(
     address : Types.Address
   ) : Result.Result<Script.Script, Text> {
-    return switch (address) {
+    switch (address) {
       case (#p2pkh p2pkhAddr) {
-        return P2pkh.makeScript(p2pkhAddr);
+        P2pkh.makeScript(p2pkhAddr);
+      };
+      case (#p2wpkh p2wpkhAddr) {
+        P2WPKH.makeScript(p2wpkhAddr);
+      };
+      case (#p2wsh p2wshAddr) {
+        P2WSH.makeScript(p2wshAddr);
       };
       case (#p2tr_key p2trKeyAddr) {
         P2tr.makeScriptFromP2trKeyAddress(p2trKeyAddr);
@@ -50,6 +82,12 @@ module {
       case (#p2pkh address1, #p2pkh address2) {
         address1 == address2;
       };
+      case (#p2wpkh addr1, #p2wpkh addr2) {
+        addr1 == addr2;
+      };
+      case (#p2wsh addr1, #p2wsh addr2) {
+        addr1 == addr2;
+      };
       case (#p2tr_key address1, #p2tr_key address2) {
         address1 == address2;
       };

src/bitcoin/P2wpkh.mo

@@ -0,0 +1,114 @@
+import Types "Types";
+import Script "./Script";
+import Segwit "../Segwit";
+import EcdsaTypes "../ecdsa/Types";
+import Result "mo:base/Result";
+import Nat8 "mo:base/Nat8";
+import Nat "mo:base/Nat";
+import Hash "../Hash";
+
+module {
+  public type DecodedP2wpkhAddress = {
+    hrp : Text;
+    publicKeyHash : [Nat8];
+  };
+
+  /// Creates the scriptPubKey for a P2WPKH address (v0).
+  /// The resulting script is: OP_0 PUSH_20 <20_byte_pubkey_hash>
+  ///
+  /// # Parameters:
+  /// - `address`: The P2WPKH address in Bech32 format (e.g., "bc1q...")
+  ///
+  /// # Returns:
+  /// `Result.Result<Script.Script, Text>` containing the Script (`#ok`) or an error message (`#err`).
+  public func makeScript(address : Types.P2WPkhAddress) : Result.Result<Script.Script, Text> {
+    switch (Segwit.decode(address)) {
+      case (#ok((_hrp, witnessProgram))) {
+        if (witnessProgram.version != 0) {
+          return #err("P2WPKH.makeScript: Invalid witness version for P2WPKH (expected 0, got " # Nat8.toText(witnessProgram.version) # ")");
+        };
+        if (witnessProgram.program.size() != 20) {
+          return #err("P2WPKH.makeScript: Invalid program size for P2WPKH (expected 20, got " # Nat.toText(witnessProgram.program.size()) # ")");
+        };
+
+        let script : Script.Script = [
+          #opcode(#OP_0),
+          #data(witnessProgram.program),
+        ];
+        #ok(script);
+      };
+      case (#err e) {
+        #err("Internal error in P2WPKH.makeScript: Failed to re-decode valid P2WPKH address '" # address # "': " # e);
+      };
+    };
+  };
+
+  func getHrp(network : Types.Network) : Text {
+    switch (network) {
+      case (#Mainnet) { "bc" };
+      case (#Testnet) { "tb" };
+      case (#Regtest) { "bcrt" }; // Common HRP for Regtest Bech32
+    };
+  };
+
+  /// Derives a P2WPKH (Bech32) address from an SEC1 public key.
+  /// Requires the public key to be in compressed format (33 bytes).
+  ///
+  /// # Parameters:
+  /// - `network`: The Bitcoin network (Mainnet, Testnet, Regtest).
+  /// - `sec1PublicKey`: The public key in SEC1 format (pair: bytes, curve). Must be compressed.
+  ///
+  /// # Returns:
+  /// `Result.Result<Types.P2wpkhAddress, Text>` containing the Bech32 address (`#ok`) or an error message (`#err`).
+  public func deriveAddress(
+    network : Types.Network,
+    sec1PublicKey : EcdsaTypes.Sec1PublicKey,
+  ) : Result.Result<Types.P2WPkhAddress, Text> {
+    let (pkBytes, _curve) = sec1PublicKey;
+
+    // P2WPKH REQUIRES a compressed public key
+    // Assuming pkBytes.size() == 33 for compressed
+    // (The library might have a helper function like PublicKey.isCompressed(pkBytes))
+    if (pkBytes.size() != 33) {
+      // Optional: We could try to compress it if it isn't, but requiring it is safer.
+      return #err("P2WPKH requires a compressed public key (33 bytes)");
+    };
+
+    // 1. Calculate HASH160(pubkey)
+    let pubKeyHash : [Nat8] = Hash.hash160(pkBytes);
+    if (pubKeyHash.size() != 20) {
+      return #err("Internal error: HASH160 result is not 20 bytes");
+    };
+
+    let hrp = getHrp(network);
+
+    let witnessProgram : Segwit.WitnessProgram = {
+      version = 0;
+      program = pubKeyHash;
+    };
+
+    return Segwit.encode(hrp, witnessProgram);
+  };
+
+  public func decodeAddress(address : Types.P2WPkhAddress) : Result.Result<DecodedP2wpkhAddress, Text> {
+    switch (Segwit.decode(address)) {
+      case (#ok((hrp, witnessProgram))) {
+        // P2WPKH specific validations
+        if (witnessProgram.version != 0) {
+          return #err("P2WPKH.decodeAddress: Invalid witness version (expected 0, got " # Nat8.toText(witnessProgram.version) # ")");
+        };
+        if (witnessProgram.program.size() != 20) {
+          return #err("P2WPKH.decodeAddress: Invalid program size (expected 20, got " # Nat.toText(witnessProgram.program.size()) # ")");
+        };
+        // Validate HRP if necessary (e.g., hrp == getHrp(#Mainnet) or hrp == getHrp(#Testnet) ...)
+
+        // Return the specific structure
+        #ok({ hrp = hrp; publicKeyHash = witnessProgram.program });
+      };
+      case (#err e) {
+        #err("P2WPKH.decodeAddress: Failed to decode Bech32 address: " # e);
+      };
+    };
+  };
+
+};

src/bitcoin/P2wsh.mo

@@ -0,0 +1,163 @@
+import Types "Types";
+import Script "./Script";
+import Segwit "../Segwit";
+import Common "../Common";
+import Result "mo:base/Result";
+import Nat8 "mo:base/Nat8";
+import Nat "mo:base/Nat";
+import Debug "mo:base/Debug";
+import Buffer "mo:base/Buffer";
+import Array "mo:base/Array";
+import Nat16 "mo:base/Nat16";
+import Nat32 "mo:base/Nat32";
+import Blob "mo:base/Blob";
+import Sha256 "mo:sha2/Sha256";
+
+module {
+
+    let opPushData1Code : Nat8 = 0x4c;
+    let opPushData1Threshold : Nat = 76;
+    let maxNat8 : Nat = 0xff;
+    let maxNat16 : Nat = 0xffff;
+
+    public type DecodedP2wshAddress = {
+        hrp : Text;
+        scriptHash : [Nat8];
+    };
+
+    func encodeOpcode(opcode : Script.Opcode) : Nat8 {
+        return switch (opcode) {
+            case (#OP_0) { 0x00 }; // for makeScript
+            case (#OP_PUSHDATA1) { opPushData1Code }; // for raw serialization
+            case (#OP_PUSHDATA2) { 0x4d }; // for raw serialization
+            case (#OP_PUSHDATA4) { 0x4e }; // for raw serialization
+            case _ {
+                Debug.trap("P2WSH internal: Unsupported opcode in encodeOpcode local replica");
+            };
+        };
+    };
+
+    func serializeScriptRaw(script : Script.Script) : [Nat8] {
+        let buf = Buffer.Buffer<Nat8>(script.size());
+
+        for (instruction in script.vals()) {
+            switch (instruction) {
+                case (#opcode(opcode)) {
+                    buf.add(encodeOpcode(opcode));
+                };
+                case (#data data) {
+                    let dataSize = data.size();
+                    if (dataSize < opPushData1Threshold) {
+                        buf.add(Nat8.fromNat(dataSize));
+                    } else if (dataSize <= maxNat8) {
+                        buf.add(encodeOpcode(#OP_PUSHDATA1));
+                        buf.add(Nat8.fromNat(dataSize));
+                    } else if (dataSize <= maxNat16) {
+                        buf.add(encodeOpcode(#OP_PUSHDATA2));
+                        let sizeData = Array.init<Nat8>(2, 0);
+                        Common.writeLE16(sizeData, 0, Nat16.fromNat(dataSize));
+                        for (byte in sizeData.vals()) { buf.add(byte) };
+                    } else {
+                        buf.add(encodeOpcode(#OP_PUSHDATA4));
+                        let sizeData = Array.init<Nat8>(4, 0);
+                        Common.writeLE32(sizeData, 0, Nat32.fromNat(dataSize));
+                        for (byte in sizeData.vals()) { buf.add(byte) };
+                    };
+                    for (byte in data.vals()) {
+                        buf.add(byte);
+                    };
+                };
+            };
+        };
+        return Buffer.toArray(buf);
+    };
+
+    /// Creates the scriptPubKey for a P2WSH address (v0).
+    /// The resulting script is:  OP_0 PUSH_32 <32_byte_script_hash>
+    ///
+    /// # Parameters:
+    /// - `address`: The P2WSH address in Bech32 format (e.g., "bc1q...")
+    ///
+    /// # Returns:
+    /// `Result.Result<Script.Script, Text>` containing the Script (`#ok`) or an error message (`#err`).
+    public func makeScript(address : Types.P2WShAddress) : Result.Result<Script.Script, Text> {
+        switch (Segwit.decode(address)) {
+            case (#ok((_hrp, witnessProgram))) {
+                if (witnessProgram.version != 0) {
+                    return #err("P2WSH.makeScript: Invalid witness version for P2WSH (expected 0, got " # Nat8.toText(witnessProgram.version) # ")");
+                };
+                if (witnessProgram.program.size() != 32) {
+                    return #err("P2WSH.makeScript: Invalid program size for P2WSH (expected 32, got " # Nat.toText(witnessProgram.program.size()) # ")");
+                };
+
+                let script : Script.Script = [
+                    #opcode(#OP_0),
+                    #data(witnessProgram.program),
+                ];
+                #ok(script);
+            };
+            case (#err e) {
+                #err("Internal error in P2WSH.makeScript: Failed to re-decode valid P2WSH address '" # address # "': " # e);
+            };
+        };
+    };
+
+    func getHrp(network : Types.Network) : Text {
+        switch (network) {
+            case (#Mainnet) { "bc" };
+            case (#Testnet) { "tb" };
+            case (#Regtest) { "bcrt" };
+        };
+    };
+
+    /// Derivate a P2WSH address (Bech32) from a witness script.
+    ///
+    /// # Parameters:
+    /// - `network`: Bitcoin network (Mainnet, Testnet, Regtest).
+    /// - `witnessScript`: the script (as `Script.Script`) wich hash SHA256 will be used.
+    ///
+    /// # Returns:
+    /// `Result.Result<Types.P2wshAddress, Text>` contains the Bech32 address (`#ok`) or an error message (`#err`).
+    public func deriveAddress(
+        network : Types.Network,
+        witnessScript : Script.Script
+    ) : Result.Result<Types.P2WShAddress, Text> {
+
+        let rawScriptBytes = serializeScriptRaw(witnessScript);
+
+        let scriptHashBlob = Sha256.fromArray(#sha256, rawScriptBytes);
+        let scriptHash = Blob.toArray(scriptHashBlob);
+
+        if (scriptHash.size() != 32) {
+            return #err("Internal error: SHA256 result is not 32 bytes");
+        };
+
+        let hrp = getHrp(network);
+
+        let witnessProgram : Segwit.WitnessProgram = {
+            version = 0;
+            program = scriptHash;
+        };
+
+        return Segwit.encode(hrp, witnessProgram);
+    };
+
+    public func decodeAddress(address : Types.P2WShAddress) : Result.Result<DecodedP2wshAddress, Text> {
+        switch (Segwit.decode(address)) {
+            case (#ok((hrp, witnessProgram))) {
+                if (witnessProgram.version != 0) {
+                    return #err("P2WSH.decodeAddress: Invalid witness version (expected 0, got " # Nat8.toText(witnessProgram.version) # ")");
+                };
+                if (witnessProgram.program.size() != 32) {
+                    return #err("P2WSH.decodeAddress: Invalid program size (expected 32, got " # Nat.toText(witnessProgram.program.size()) # ")");
+                };
+
+                #ok({ hrp = hrp; scriptHash = witnessProgram.program });
+            };
+            case (#err e) {
+                #err("P2WSH.decodeAddress: Failed to decode Bech32 address: " # e);
+            };
+        };
+    };
+
+};