kabu_execution_multicaller/
opcodes_encoder.rs

1use std::ops::{BitOrAssign, Shl};
2
3use alloy_primitives::{Bytes, U256};
4use alloy_sol_types::SolInterface;
5use eyre::{ErrReport, Result};
6use lazy_static::lazy_static;
7
8use kabu_defi_abi::multicaller::IMultiCaller;
9use kabu_types_blockchain::{CallType, MulticallerCall, MulticallerCalls};
10
11lazy_static! {
12    static ref VALUE_CALL_SELECTOR: U256 = U256::from(0x7FFA);
13    static ref CALCULATION_CALL_SELECTOR: U256 = U256::from(0x7FFB);
14    static ref ZERO_VALUE_CALL_SELECTOR: U256 = U256::from(0x7FFC);
15    static ref INTERNAL_CALL_SELECTOR: U256 = U256::from(0x7FFD);
16    static ref STATIC_CALL_SELECTOR: U256 = U256::from(0x7FFE);
17    static ref DELEGATE_CALL_SELECTOR: U256 = U256::from(0x7FFF);
18}
19
20pub struct OpcodesEncoderV2;
21
22pub trait OpcodesEncoder {
23    fn pack_do_calls(opcodes: &MulticallerCalls) -> Result<Bytes>;
24    fn pack_do_calls_data(opcode: &MulticallerCalls) -> Result<Bytes>;
25}
26
27impl OpcodesEncoderV2 {
28    fn encode_data_offset(is_relative: bool, stack_offset: u32, data_offset: u32, data_len: usize) -> u32 {
29        let mut ret = if is_relative { 0x800000 } else { 0x0 };
30        ret |= (stack_offset & 0x7) << 20;
31        ret |= (data_len as u32 & 0xFF) << 12;
32        ret |= data_offset & 0xFFF;
33        ret
34    }
35
36    pub fn encode_call_stack(opcode: &MulticallerCall) -> u32 {
37        if let Some(call_stack) = &opcode.call_stack {
38            match opcode.call_type {
39                CallType::InternalCall | CallType::CalculationCall => Self::encode_data_offset(
40                    call_stack.is_relative,
41                    call_stack.stack_offset,
42                    call_stack.data_offset + 0xC,
43                    call_stack.data_len,
44                ),
45                _ => Self::encode_data_offset(
46                    call_stack.is_relative,
47                    call_stack.stack_offset,
48                    call_stack.data_offset + 0x20,
49                    call_stack.data_len,
50                ),
51            }
52        } else {
53            0xFFFFFF
54        }
55    }
56
57    pub fn encode_return_stack(opcode: &MulticallerCall) -> u32 {
58        if let Some(return_stack) = &opcode.return_stack {
59            Self::encode_data_offset(return_stack.is_relative, return_stack.stack_offset, return_stack.data_offset, return_stack.data_len)
60        } else {
61            0xFFFFFF
62        }
63    }
64
65    fn pack_opcode(opcode: &MulticallerCall) -> Result<Vec<u8>> {
66        let mut ret: Vec<u8> = Vec::new();
67        let mut selector = U256::ZERO;
68        //let mut selector_bytes_len = 0x20;
69        let selector_call = match opcode.call_type {
70            CallType::Call => {
71                if opcode.value.is_none() {
72                    *ZERO_VALUE_CALL_SELECTOR
73                } else {
74                    *VALUE_CALL_SELECTOR
75                }
76            }
77            CallType::DelegateCall => *DELEGATE_CALL_SELECTOR,
78            CallType::InternalCall => {
79                //selector_bytes_len = 0xC;
80                *INTERNAL_CALL_SELECTOR
81            }
82            CallType::StaticCall => *STATIC_CALL_SELECTOR,
83            CallType::CalculationCall => {
84                //selector_bytes_len = 0xC;
85                *CALCULATION_CALL_SELECTOR
86            }
87            _ => {
88                return Err(ErrReport::msg("WRONG_OPCODE"));
89            }
90        };
91
92        if selector_call == *VALUE_CALL_SELECTOR && !opcode.value.unwrap_or_default().is_zero() {
93            selector = opcode.value.unwrap_or_default().shl(0x10);
94            selector.bitor_assign(U256::from(1).shl(96 - 1));
95            selector.bitor_assign(U256::from(opcode.call_data.len()).shl(0));
96        } else {
97            selector.bitor_assign(selector_call.shl(80));
98            selector.bitor_assign(U256::from(opcode.call_data.len()).shl(0));
99            selector.bitor_assign(U256::from(Self::encode_call_stack(opcode)).shl(16));
100            selector.bitor_assign(U256::from(Self::encode_return_stack(opcode)).shl(40));
101        }
102
103        let selector_bytes = selector.to_be_bytes::<32>();
104        ret.append(&mut selector_bytes[20..32].to_vec());
105
106        match opcode.call_type {
107            CallType::CalculationCall | CallType::InternalCall => {}
108            _ => {
109                let mut address_bytes = opcode.to.to_vec();
110                ret.append(&mut address_bytes);
111            }
112        }
113
114        ret.append(&mut opcode.call_data.to_vec());
115
116        Ok(ret)
117    }
118}
119
120impl OpcodesEncoder for OpcodesEncoderV2 {
121    fn pack_do_calls(opcodes: &MulticallerCalls) -> Result<Bytes> {
122        let call_data = OpcodesEncoderV2::pack_do_calls_data(opcodes)?;
123        let args = IMultiCaller::doCallsCall { data: call_data };
124        let call = IMultiCaller::IMultiCallerCalls::doCalls(args);
125        Ok(call.abi_encode().into())
126    }
127
128    fn pack_do_calls_data(opcodes: &MulticallerCalls) -> Result<Bytes> {
129        let mut call_data: Vec<u8> = Vec::new();
130        for o in opcodes.opcodes_vec.iter() {
131            call_data.append(&mut OpcodesEncoderV2::pack_opcode(o)?);
132        }
133        Ok(call_data.into())
134    }
135}
136
137#[cfg(test)]
138mod test {
139    use super::*;
140
141    #[test]
142    fn test() {
143        let buf = Bytes::from(vec![0x33, 0x33, 0x44, 0x55]);
144
145        let mut opcode = MulticallerCall::new_internal_call(&buf);
146        //let mut opcode = Opcode::new_internal_call(to, &Some(buf));
147        opcode.set_call_stack(true, 0, 24, 0x20).set_return_stack(true, 1, 44, 0x20);
148
149        let mut opcodes = MulticallerCalls::new();
150        opcodes.add(opcode);
151
152        let packed_bytes = OpcodesEncoderV2::pack_do_calls(&opcodes).unwrap();
153        println!("{packed_bytes:?}");
154    }
155}