core: Add Affinity bytecode

Apply affinities to a range of P2 registers starting with P1.

P4 is a string that is P2 characters long. The N-th character of the string indicates the column affinity that should be used for the N-th memory cell in the range.

core/schema.rs

@@ -1,5 +1,5 @@
-use crate::VirtualTable;
 use crate::{util::normalize_ident, Result};
+use crate::{LimboError, VirtualTable};
 use core::fmt;
 use fallible_iterator::FallibleIterator;
 use limbo_sqlite3_parser::ast::{Expr, Literal, SortOrder, TableOptions};
@@ -585,6 +585,20 @@ impl Affinity {
             Affinity::Numeric => SQLITE_AFF_NUMERIC,
         }
     }
+
+    pub fn from_char(char: char) -> Result<Self> {
+        match char {
+            SQLITE_AFF_INTEGER => Ok(Affinity::Integer),
+            SQLITE_AFF_TEXT => Ok(Affinity::Text),
+            SQLITE_AFF_NONE => Ok(Affinity::Blob),
+            SQLITE_AFF_REAL => Ok(Affinity::Real),
+            SQLITE_AFF_NUMERIC => Ok(Affinity::Numeric),
+            _ => Err(LimboError::InternalError(format!(
+                "Invalid affinity character: {}",
+                char
+            ))),
+        }
+    }
 }
 
 impl fmt::Display for Type {

core/vdbe/execute.rs

@@ -4509,7 +4509,7 @@ pub fn op_not_found(
             return_if_io!(cursor.seek(SeekKey::IndexKey(&record), SeekOp::EQ))
         }
     };
-    
+
     if found {
         state.pc += 1;
     } else {
@@ -4519,6 +4519,40 @@ pub fn op_not_found(
     Ok(InsnFunctionStepResult::Step)
 }
 
+pub fn op_affinity(
+    program: &Program,
+    state: &mut ProgramState,
+    insn: &Insn,
+    pager: &Rc<Pager>,
+    mv_store: Option<&Rc<MvStore>>,
+) -> Result<InsnFunctionStepResult> {
+    let Insn::Affinity {
+        start_reg,
+        count,
+        affinities,
+    } = insn
+    else {
+        unreachable!("unexpected Insn {:?}", insn)
+    };
+
+    if affinities.len() != count.get() {
+        return Err(LimboError::InternalError(
+            "Affinity: the length of affinities does not match the count".into(),
+        ));
+    }
+
+    for (i, affinity_char) in affinities.chars().enumerate().take(count.get()) {
+        let reg_index = *start_reg + i;
+
+        let affinity = Affinity::from_char(affinity_char)?;
+
+        apply_affinity_char(&mut state.registers[reg_index], affinity);
+    }
+
+    state.pc += 1;
+    Ok(InsnFunctionStepResult::Step)
+}
+
 fn exec_lower(reg: &OwnedValue) -> Option<OwnedValue> {
     match reg {
         OwnedValue::Text(t) => Some(OwnedValue::build_text(&t.as_str().to_lowercase())),

core/vdbe/explain.rs

@@ -1379,6 +1379,28 @@ pub fn insn_to_str(
                     target_pc.to_debug_int()
                 ),
             ),
+            Insn::Affinity {
+                start_reg,
+                count,
+                affinities,
+            } => (
+                "Affinity",
+                *start_reg as i32,
+                count.get() as i32,
+                0,
+                OwnedValue::build_text(""),
+                0,
+                format!(
+                    "r[{}..{}] = {}",
+                    start_reg,
+                    start_reg + count.get(),
+                    affinities
+                        .chars()
+                        .map(|a| a.to_string())
+                        .collect::<Vec<_>>()
+                        .join(", ")
+                ),
+            ),
         };
     format!(
         "{:<4}  {:<17}  {:<4}  {:<4}  {:<4}  {:<13}  {:<2}  {}",

core/vdbe/insn.rs

@@ -1,4 +1,7 @@
-use std::{num::NonZero, rc::Rc};
+use std::{
+    num::{NonZero, NonZeroUsize},
+    rc::Rc,
+};
 
 use super::{execute, AggFunc, BranchOffset, CursorID, FuncCtx, InsnFunction, PageIdx};
 use crate::{
@@ -809,6 +812,12 @@ pub enum Insn {
         record_reg: usize,
         num_regs: usize,
     },
+    /// Apply affinities to a range of registers. Affinities must have the same size of count
+    Affinity {
+        start_reg: usize,
+        count: NonZeroUsize,
+        affinities: String,
+    },
 }
 
 impl Insn {
@@ -924,6 +933,7 @@ impl Insn {
             Insn::OpenEphemeral { .. } | Insn::OpenAutoindex { .. } => execute::op_open_ephemeral,
             Insn::Once { .. } => execute::op_once,
             Insn::NotFound { .. } => execute::op_not_found,
+            Insn::Affinity { .. } => execute::op_affinity,
         }
     }
 }