Lex fixup

Add build_scriptint from rust-bitcoin

Author: sanket1729

src/miniscript/decode.rs

@@ -282,7 +282,7 @@ pub fn parse<Ctx: ScriptContext>(
                                     ))?
                                 },
                             ),
-                            Tk::Push4(ver), Tk::Pick, Tk::Sub, Tk::Depth => match_token!(
+                            Tk::PickPush4(ver), Tk::Sub, Tk::Depth => match_token!(
                                 tokens,
                                 Tk::Num(2) => {
                                     non_term.push(NonTerm::Verify);
@@ -360,7 +360,7 @@ pub fn parse<Ctx: ScriptContext>(
                                 hash160::Hash::from_inner(hash)
                             ))?,
                         ),
-                        Tk::Push4(ver), Tk::Pick, Tk::Sub, Tk::Depth => match_token!(
+                        Tk::PickPush4(ver), Tk::Sub, Tk::Depth => match_token!(
                             tokens,
                             Tk::Num(2) => term.reduce0(Terminal::Version(ver))?,
                         ),

src/miniscript/lex.rs

@@ -23,7 +23,7 @@ use elements::{opcodes, script};
 use std::fmt;
 
 use super::Error;
-use util::slice_to_u32_le;
+use util::{build_scriptint, slice_to_u32_le};
 /// Atom of a tokenized version of a script
 #[derive(Debug, Clone, PartialEq, Eq)]
 #[allow(missing_docs)]
@@ -65,8 +65,10 @@ pub enum Token {
     Hash20([u8; 20]),
     Hash32([u8; 32]),
     Pubkey(PublicKey),
-    Push4(u32),
-    Push(Vec<u8>), // <pref> <swap> <cat> /* prefix is catted */
+    Push(Vec<u8>),        // Num or a
+    PickPush4(u32),       // Pick followed by a 4 byte push
+    PickPush32([u8; 32]), // Pick followed by a 32 byte push
+    PickPush(Vec<u8>),    // Pick followed by a push
 }
 
 impl fmt::Display for Token {
@@ -131,6 +133,23 @@ impl Iterator for TokenIter {
 pub fn lex(script: &script::Script) -> Result<Vec<Token>, Error> {
     let mut ret = Vec::with_capacity(script.len());
 
+    fn process_candidate_push(ret: &mut Vec<Token>) -> Result<(), Error> {
+        let ret_len = ret.len();
+
+        if ret_len < 2 || ret[ret_len - 1] != Token::Swap {
+            return Ok(());
+        }
+        let token = match &ret[ret_len - 2] {
+            Token::Hash20(x) => Token::Push(x.to_vec()),
+            Token::Hash32(x) => Token::Push(x.to_vec()),
+            Token::Pubkey(pk) => Token::Push(pk.to_bytes()),
+            Token::Num(k) => Token::Push(build_scriptint(*k as i64)),
+            _x => return Ok(()), // no change required
+        };
+        ret[ret_len - 2] = token;
+        Ok(())
+    }
+
     for ins in script.instructions_minimal() {
         match ins.map_err(Error::Script)? {
             script::Instruction::Op(opcodes::all::OP_BOOLAND) => {
@@ -179,6 +198,7 @@ pub fn lex(script: &script::Script) -> Result<Vec<Token>, Error> {
                 ret.push(Token::Left);
             }
             script::Instruction::Op(opcodes::all::OP_CAT) => {
+                process_candidate_push(&mut ret)?;
                 ret.push(Token::Cat);
             }
             script::Instruction::Op(opcodes::all::OP_CODESEPARATOR) => {
@@ -253,36 +273,41 @@ pub fn lex(script: &script::Script) -> Result<Vec<Token>, Error> {
                 ret.push(Token::Hash256);
             }
             script::Instruction::PushBytes(bytes) => {
-                // Check for Push
-                let ret_len = ret.len();
-                // See that last five elements are CAT
-                // Checking only two here
-                if ret_len >= 5
-                    && ret.last() == Some(&Token::Cat)
-                    && ret.get(ret_len - 2) == Some(&Token::Cat)
-                {
-                    ret.push(Token::Push(bytes.to_owned()));
-                }
+                // Check for Pick Push
                 // Special handling of tokens for Covenants
                 // To determine whether some Token is actually
                 // 4 bytes push or a script int of 4 bytes,
                 // we need additional script context
-                else if ret.last() == Some(&Token::Pick) {
+                if ret.last() == Some(&Token::Pick) {
+                    ret.pop().unwrap();
                     match bytes.len() {
-                        4 => ret.push(Token::Push4(slice_to_u32_le(bytes))),
+                        // All other sighash elements are 32 bytes. And the script code
+                        // is 24 bytes
+                        4 => ret.push(Token::PickPush4(slice_to_u32_le(bytes))),
+                        32 => {
+                            let mut x = [0u8; 32];
+                            x.copy_from_slice(bytes);
+                            ret.push(Token::PickPush32(x));
+                        }
+                        // Other pushes should be err. This will change
+                        // once we add script introspection
                         _ => return Err(Error::InvalidPush(bytes.to_owned())),
                     }
                 } else {
+                    // Create the most specific type possible out of the
+                    // Push. When we later encounter CAT, revisit and
+                    // reconvert these to pushes.
+                    // See [process_candidate_push]
                     match bytes.len() {
                         20 => {
                             let mut x = [0; 20];
                             x.copy_from_slice(bytes);
-                            ret.push(Token::Hash20(x))
+                            ret.push(Token::Hash20(x));
                         }
                         32 => {
                             let mut x = [0; 32];
                             x.copy_from_slice(bytes);
-                            ret.push(Token::Hash32(x))
+                            ret.push(Token::Hash32(x));
                         }
                         33 | 65 => {
                             ret.push(Token::Pubkey(
@@ -300,8 +325,7 @@ pub fn lex(script: &script::Script) -> Result<Vec<Token>, Error> {
                                     }
                                     ret.push(Token::Num(v as u32));
                                 }
-                                Ok(_) => return Err(Error::InvalidPush(bytes.to_owned())),
-                                Err(e) => return Err(Error::Script(e)),
+                                _ => ret.push(Token::Push(bytes.to_owned())),
                             }
                         }
                     }

src/util.rs

@@ -37,6 +37,34 @@ macro_rules! define_slice_to_le {
 
 define_slice_to_le!(slice_to_u32_le, u32);
 
+/// Helper to encode an integer in script format
+/// Copied from rust-bitcoin
+pub(crate) fn build_scriptint(n: i64) -> Vec<u8> {
+    if n == 0 {
+        return vec![];
+    }
+
+    let neg = n < 0;
+
+    let mut abs = if neg { -n } else { n } as usize;
+    let mut v = vec![];
+    while abs > 0xFF {
+        v.push((abs & 0xFF) as u8);
+        abs >>= 8;
+    }
+    // If the number's value causes the sign bit to be set, we need an extra
+    // byte to get the correct value and correct sign bit
+    if abs & 0x80 != 0 {
+        v.push(abs as u8);
+        v.push(if neg { 0x80u8 } else { 0u8 });
+    }
+    // Otherwise we just set the sign bit ourselves
+    else {
+        abs |= if neg { 0x80 } else { 0 };
+        v.push(abs as u8);
+    }
+    v
+}
 /// Get the count of non-push opcodes
 // Export to upstream
 #[cfg(test)]