refactor: address TODOs ExtHandle usage and simplify DslIr Ext field opcodes

crates/recursion/compiler/src/ir/arithmetic.rs

@@ -1,6 +1,6 @@
 use std::{cell::UnsafeCell, mem::ManuallyDrop};
 
-use p3_field::{AbstractExtensionField, AbstractField, Field};
+use p3_field::{AbstractExtensionField, Field};
 
 use crate::ir::DslIr;
 
@@ -455,7 +455,8 @@ impl<C: Config> ExtOperations<C::F, C::EF> for UnsafeCell<InnerBuilder<C>> {
         let mut inner = unsafe { ManuallyDrop::new(Box::from_raw(ptr as *mut Self)) };
         let inner = inner.get_mut();
         let idx = inner.variable_count;
-        let res = Ext::new(idx, lhs.handle);
+        let handle = inner.ext_handle_ptr as *mut ExtHandle<_, _>;
+        let res = Ext::new(idx, handle);
         inner.variable_count += 1;
 
         inner.operations.push(DslIr::AddE(res, lhs, rhs));
@@ -467,7 +468,7 @@ impl<C: Config> ExtOperations<C::F, C::EF> for UnsafeCell<InnerBuilder<C>> {
         let mut inner = unsafe { ManuallyDrop::new(Box::from_raw(ptr as *mut Self)) };
         let inner = inner.get_mut();
         let idx = inner.variable_count;
-        let res = Ext::new(idx, lhs.handle);
+        let res = Ext::new(idx, inner.ext_handle_ptr);
         inner.variable_count += 1;
 
         inner.operations.push(DslIr::AddEF(res, lhs, rhs));
@@ -574,9 +575,14 @@ impl<C: Config> ExtOperations<C::F, C::EF> for UnsafeCell<InnerBuilder<C>> {
     }
 
     fn sub_base_ext(ptr: *mut (), lhs: Felt<C::F>, rhs: Ext<C::F, C::EF>) -> Ext<C::F, C::EF> {
-        // TODO: optimize to one opcode.
-        let rhs = Self::neg_ext(ptr, rhs);
-        Self::add_ext_base(ptr, rhs, lhs)
+        let mut inner = unsafe { ManuallyDrop::new(Box::from_raw(ptr as *mut Self)) };
+        let inner = inner.get_mut();
+        let idx = inner.variable_count;
+        let res = Ext::new(idx, rhs.handle);
+        inner.variable_count += 1;
+
+        inner.operations.push(DslIr::SubFE(res, lhs, rhs));
+        res
     }
 
     fn neg_ext(ptr: *mut (), lhs: Ext<C::F, C::EF>) -> Ext<C::F, C::EF> {
@@ -631,11 +637,16 @@ impl<C: Config> ExtOperations<C::F, C::EF> for UnsafeCell<InnerBuilder<C>> {
         ptr: *mut (),
         lhs: Felt<C::F>,
         rhs: C::EF,
-        handle: *mut ExtHandle<C::F, C::EF>,
+        _handle: *mut ExtHandle<C::F, C::EF>,
     ) -> Ext<C::F, C::EF> {
-        // TODO: optimize to one opcode.
-        let lhs = Self::add_felt_const_ext(ptr, lhs, C::EF::zero(), handle);
-        Self::mul_const_ext(ptr, lhs, rhs)
+        let mut inner = unsafe { ManuallyDrop::new(Box::from_raw(ptr as *mut Self)) };
+        let inner = inner.get_mut();
+        let idx = inner.variable_count;
+        let res = Ext::new(idx, inner.ext_handle_ptr);
+        inner.variable_count += 1;
+
+        inner.operations.push(DslIr::MulFEI(res, lhs, rhs));
+        res
     }
 
     fn div_ext(ptr: *mut (), lhs: Ext<C::F, C::EF>, rhs: Ext<C::F, C::EF>) -> Ext<C::F, C::EF> {
@@ -667,9 +678,14 @@ impl<C: Config> ExtOperations<C::F, C::EF> for UnsafeCell<InnerBuilder<C>> {
     }
 
     fn div_base_ext(ptr: *mut (), lhs: Felt<C::F>, rhs: Ext<C::F, C::EF>) -> Ext<C::F, C::EF> {
-        // TODO: optimize to one opcode.
-        let lhs = Self::add_felt_const_ext(ptr, lhs, C::EF::zero(), rhs.handle);
-        Self::div_ext(ptr, lhs, rhs)
+        let mut inner = unsafe { ManuallyDrop::new(Box::from_raw(ptr as *mut Self)) };
+        let inner = inner.get_mut();
+        let idx = inner.variable_count;
+        let res = Ext::new(idx, rhs.handle);
+        inner.variable_count += 1;
+
+        inner.operations.push(DslIr::DivFE(res, lhs, rhs));
+        res
     }
 
     fn div_ext_base(ptr: *mut (), lhs: Ext<C::F, C::EF>, rhs: Felt<C::F>) -> Ext<C::F, C::EF> {

crates/recursion/compiler/src/ir/builder.rs

@@ -13,6 +13,7 @@ use super::{
 pub struct InnerBuilder<C: Config> {
     pub(crate) variable_count: u32,
     pub operations: Vec<DslIr<C>>,
+    pub ext_handle_ptr: *mut ExtHandle<C::F, C::EF>,
 }
 
 /// A builder for the DSL.
@@ -48,15 +49,20 @@ impl<C: Config> Builder<C> {
         let mut inner = Box::new(UnsafeCell::new(InnerBuilder {
             variable_count: 0,
             operations: Default::default(),
+            ext_handle_ptr: std::ptr::null_mut(),
         }));
 
-        let var_handle = Box::new(VarOperations::var_handle(&mut inner));
         let mut ext_handle = Box::new(ExtOperations::ext_handle(&mut inner));
+
+        inner.get_mut().ext_handle_ptr = ext_handle.as_mut() as *mut _;
+
         let felt_handle = Box::new(FeltOperations::felt_handle(
             &mut inner,
             ext_handle.as_mut() as *mut _ as *mut (),
         ));
 
+        let var_handle = Box::new(VarOperations::var_handle(&mut inner));
+
         let mut new_builder = Self {
             inner,
             witness_var_count: 0,

crates/recursion/compiler/src/ir/instructions.rs

@@ -58,6 +58,8 @@ pub enum DslIr<C: Config> {
     SubF(Felt<C::F>, Felt<C::F>, Felt<C::F>),
     /// Subtracts a field element and a field immediate (felt = felt - field imm).
     SubFI(Felt<C::F>, Felt<C::F>, C::F),
+    /// Subtracts a field element and an extension field element (ext = felt - ext).
+    SubFE(Ext<C::F, C::EF>, Felt<C::F>, Ext<C::F, C::EF>),
     /// Subtracts a field immediate and a field element (felt = field imm - felt).
     SubFIN(Felt<C::F>, C::F, Felt<C::F>),
     /// Subtracts two extension field elements (ext = ext - ext).
@@ -80,6 +82,8 @@ pub enum DslIr<C: Config> {
     MulVI(Var<C::N>, Var<C::N>, C::N),
     /// Multiplies two field elements (felt = felt * felt).
     MulF(Felt<C::F>, Felt<C::F>, Felt<C::F>),
+    /// Multiplies a felt and an extension field immediate (ext = felt * ext field imm).
+    MulFEI(Ext<C::F, C::EF>, Felt<C::F>, C::EF),
     /// Multiplies a field element and a field immediate (felt = felt * field imm).
     MulFI(Felt<C::F>, Felt<C::F>, C::F),
     /// Multiplies two extension field elements (ext = ext * ext).
@@ -95,6 +99,8 @@ pub enum DslIr<C: Config> {
     // Divisions.
     /// Divides two variables (var = var / var).
     DivF(Felt<C::F>, Felt<C::F>, Felt<C::F>),
+    /// Divides a field element and an extension field element (ext = felt / ext).
+    DivFE(Ext<C::F, C::EF>, Felt<C::F>, Ext<C::F, C::EF>),
     /// Divides a field element and a field immediate (felt = felt / field imm).
     DivFI(Felt<C::F>, Felt<C::F>, C::F),
     /// Divides a field immediate and a field element (felt = field imm / felt).