Implement Interpret for Gep

crates/codegen/src/optim/sccp.rs

@@ -13,7 +13,7 @@ use sonatina_ir::{
     inst::control_flow::{Branch, BranchKind},
     interpret::{Action, EvalValue, Interpret, State},
     prelude::*,
-    BlockId, ControlFlowGraph, Function, Immediate, InstId, Type, ValueId,
+    BlockId, ControlFlowGraph, DataFlowGraph, Function, Immediate, InstId, Type, ValueId,
 };
 
 #[derive(Debug)]
@@ -180,7 +180,7 @@ impl SccpSolver {
             return;
         }
 
-        let mut cell_state = CellState::new(&self.lattice);
+        let mut cell_state = CellState::new(&self.lattice, &func.dfg);
         let value = InstDowncast::map(func.inst_set(), inst, |i: &dyn Interpret| {
             i.interpret(&mut cell_state)
         });
@@ -456,15 +456,17 @@ impl Default for LatticeCell {
     }
 }
 
-struct CellState<'i> {
+struct CellState<'a, 'i> {
     map: &'i SecondaryMap<ValueId, LatticeCell>,
     used_val: FxHashSet<ValueId>,
+    dfg: &'a DataFlowGraph,
 }
-impl<'i> CellState<'i> {
-    fn new(map: &'i SecondaryMap<ValueId, LatticeCell>) -> Self {
+impl<'a, 'i> CellState<'a, 'i> {
+    fn new(map: &'i SecondaryMap<ValueId, LatticeCell>, dfg: &'a DataFlowGraph) -> Self {
         Self {
             map,
             used_val: Default::default(),
+            dfg,
         }
     }
 
@@ -479,7 +481,7 @@ impl<'i> CellState<'i> {
     }
 }
 
-impl<'i> State for CellState<'i> {
+impl<'a, 'i> State for CellState<'a, 'i> {
     fn lookup_val(&mut self, value: ValueId) -> EvalValue {
         self.used_val.insert(value);
 
@@ -514,4 +516,8 @@ impl<'i> State for CellState<'i> {
     fn store(&mut self, _addr: EvalValue, _value: EvalValue, _ty: Type) -> EvalValue {
         panic!("instruction with side effect must not be interpreted")
     }
+
+    fn dfg(&self) -> &DataFlowGraph {
+        self.dfg
+    }
 }

crates/ir/src/interpret/arith.rs

@@ -35,7 +35,7 @@ impl Interpret for Mul {
         let rhs = state.lookup_val(*self.rhs());
         state.set_action(Action::Continue);
 
-        EvalValue::zip_with_imm(lhs, rhs, |lhs, rhs| lhs - rhs)
+        EvalValue::zip_with_imm(lhs, rhs, |lhs, rhs| lhs * rhs)
     }
 }
 

crates/ir/src/interpret/control_flow.rs

@@ -57,6 +57,7 @@ impl Interpret for BrTable {
 impl Interpret for Phi {
     fn interpret(&self, state: &mut dyn State) -> EvalValue {
         let prev_block = state.prev_block();
+        state.set_action(Action::Continue);
         for (value, block) in self.args() {
             if prev_block == *block {
                 return state.lookup_val(*value);

crates/ir/src/interpret/data.rs

@@ -1,22 +1,100 @@
 use super::{Action, EvalValue, Interpret, State};
-use crate::inst::data::*;
+use crate::{inst::data::*, types::CompoundTypeData, Immediate, Type, I256};
 
 impl Interpret for Mload {
     fn interpret(&self, state: &mut dyn State) -> EvalValue {
-        let addr = state.lookup_val(*self.addr());
         state.set_action(Action::Continue);
 
+        let addr = state.lookup_val(*self.addr());
         state.load(addr, *self.ty())
     }
 }
 
 impl Interpret for Mstore {
     fn interpret(&self, state: &mut dyn State) -> EvalValue {
+        state.set_action(Action::Continue);
+
         let value = state.lookup_val(*self.addr());
         let addr = state.lookup_val(*self.addr());
+        state.store(addr, value, *self.ty())
+    }
+}
 
+impl Interpret for Gep {
+    fn interpret(&self, state: &mut dyn State) -> EvalValue {
         state.set_action(Action::Continue);
 
-        state.store(addr, value, *self.ty())
+        let Some(base_addr) = state.lookup_val(self.values()[0]).as_imm() else {
+            return EvalValue::Undef;
+        };
+
+        let mut current_ty = state.dfg().value_ty(self.values()[0]);
+
+        assert!(
+            state.dfg().ctx.with_ty_store(|s| s.is_ptr(current_ty)),
+            "GEP must start with a pointer type"
+        );
+
+        let mut offset = 0;
+        for value in self.values()[1..].iter() {
+            let Some(idx_value) = state.lookup_val(*value).as_imm() else {
+                return EvalValue::Undef;
+            };
+            let idx_value = idx_value.as_usize();
+
+            let cmpd = match current_ty {
+                Type::I1
+                | Type::I16
+                | Type::I8
+                | Type::I32
+                | Type::I64
+                | Type::I128
+                | Type::I256
+                | Type::Unit => {
+                    panic!("Invalid GEP: indexing into a scalar type or unit with more indices remaining");
+                }
+                Type::Compound(cmpd) => cmpd,
+            };
+
+            let cmpd_data = state
+                .dfg()
+                .ctx
+                .with_ty_store(|s| s.resolve_compound(cmpd).clone());
+
+            match cmpd_data {
+                CompoundTypeData::Array { elem, .. } => {
+                    let elem_size = state.dfg().ctx.size_of(elem);
+                    offset += elem_size * idx_value;
+                    current_ty = elem;
+                }
+
+                CompoundTypeData::Ptr(ty) => {
+                    let size = state.dfg().ctx.size_of(ty);
+                    offset += size * idx_value;
+                    current_ty = ty;
+                }
+
+                CompoundTypeData::Struct(s) => {
+                    let mut local_offset = 0;
+                    for i in 0..idx_value.saturating_sub(1) {
+                        let field_ty = s.fields[i];
+                        let size = state.dfg().ctx.size_of(field_ty);
+                        let align = state.dfg().ctx.align_of(field_ty);
+                        local_offset += align_to(offset + size, align);
+                    }
+                    offset += local_offset;
+                    current_ty = s.fields[idx_value];
+                }
+            }
+        }
+
+        let tl = state.dfg().ctx.type_layout;
+        let res = base_addr + Immediate::from_i256(I256::from(offset), tl.pointer_repl());
+        EvalValue::Imm(res)
     }
 }
+
+fn align_to(offset: usize, alignment: usize) -> usize {
+    assert!(alignment & (alignment - 1) == 0);
+    (offset + alignment - 1) & !(alignment - 1)
+}

crates/ir/src/interpret/mod.rs

@@ -1,6 +1,6 @@
 use macros::inst_prop;
 
-use crate::{inst, module::FuncRef, BlockId, Immediate, Type, ValueId};
+use crate::{inst, module::FuncRef, BlockId, DataFlowGraph, Immediate, Type, ValueId};
 
 mod arith;
 mod cast;
@@ -13,6 +13,8 @@ mod logic;
 pub trait Interpret {
     fn interpret(&self, state: &mut dyn State) -> EvalValue;
 
+    // TODO: Implement `Interpret` for all inst types and use
+    // `type Members = All`
     type Members = (
         inst::arith::Neg,
         inst::arith::Add,
@@ -42,6 +44,7 @@ pub trait Interpret {
         inst::cmp::IsZero,
         inst::data::Mload,
         inst::data::Mstore,
+        inst::data::Gep,
         inst::control_flow::Jump,
         inst::control_flow::Br,
         inst::control_flow::BrTable,
@@ -82,6 +85,8 @@ pub trait State {
     fn load(&mut self, addr: EvalValue, ty: Type) -> EvalValue;
 
     fn store(&mut self, addr: EvalValue, value: EvalValue, ty: Type) -> EvalValue;
+
+    fn dfg(&self) -> &DataFlowGraph;
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]

crates/ir/src/isa/evm.rs

@@ -3,7 +3,7 @@ use std::sync::LazyLock;
 use sonatina_triple::{Architecture, TargetTriple};
 
 use super::{Endian, Isa, TypeLayout};
-use crate::{inst::evm::inst_set::EvmInstSet, types::CompoundTypeData, Type};
+use crate::{inst::evm::inst_set::EvmInstSet, module::ModuleCtx, types::CompoundTypeData, Type};
 
 #[derive(Debug, Clone, Copy)]
 pub struct Evm {
@@ -37,14 +37,14 @@ impl Isa for Evm {
 
 struct EvmTypeLayout {}
 impl TypeLayout for EvmTypeLayout {
-    fn size_of(&self, ty: crate::Type, ty_store: &crate::types::TypeStore) -> usize {
+    fn size_of(&self, ty: crate::Type, ctx: &ModuleCtx) -> usize {
         match ty {
             Type::Compound(cmpd) => {
-                let cmpd_data = ty_store.resolve_compound(cmpd);
+                let cmpd_data = ctx.with_ty_store(|s| s.resolve_compound(cmpd).clone());
                 match cmpd_data {
                     CompoundTypeData::Array { elem, len } => {
                         // TODO: alignment!
-                        self.size_of(*elem, ty_store) * len
+                        self.size_of(elem, ctx) * len
                     }
 
                     CompoundTypeData::Ptr(_) => 32,
@@ -56,7 +56,7 @@ impl TypeLayout for EvmTypeLayout {
                         s.fields
                             .iter()
                             .copied()
-                            .fold(0, |acc, ty| acc + self.size_of(ty, ty_store))
+                            .fold(0, |acc, ty| acc + self.size_of(ty, ctx))
                     }
                 }
             }
@@ -67,6 +67,14 @@ impl TypeLayout for EvmTypeLayout {
         }
     }
 
+    fn pointer_repl(&self) -> Type {
+        Type::I256
+    }
+
+    fn align_of(&self, _ty: Type, _ctx: &ModuleCtx) -> usize {
+        1
+    }
+
     fn endian(&self) -> Endian {
         Endian::Be
     }

crates/ir/src/isa/mod.rs

@@ -1,6 +1,6 @@
 use sonatina_triple::TargetTriple;
 
-use crate::{types::TypeStore, InstSetBase, Type};
+use crate::{module::ModuleCtx, InstSetBase, Type};
 
 pub mod evm;
 
@@ -13,8 +13,9 @@ pub trait Isa {
 }
 
 pub trait TypeLayout {
-    fn size_of(&self, ty: Type, ty_store: &TypeStore) -> usize;
-
+    fn size_of(&self, ty: Type, ctx: &ModuleCtx) -> usize;
+    fn align_of(&self, ty: Type, ctx: &ModuleCtx) -> usize;
+    fn pointer_repl(&self) -> Type;
     fn endian(&self) -> Endian;
 }
 

crates/ir/src/module.rs

@@ -59,7 +59,11 @@ impl ModuleCtx {
     }
 
     pub fn size_of(&self, ty: Type) -> usize {
-        self.with_ty_store(|ty_store| self.type_layout.size_of(ty, ty_store))
+        self.type_layout.size_of(ty, self)
+    }
+
+    pub fn align_of(&self, ty: Type) -> usize {
+        self.type_layout.align_of(ty, self)
     }
 
     pub fn func_sig(&self, func: FuncRef) -> Option<&Signature> {

crates/ir/src/value.rs

@@ -122,12 +122,12 @@ impl Immediate {
     }
 
     pub fn sext(self, ty: Type) -> Self {
-        debug_assert!(self.ty() < ty);
+        debug_assert!(self.ty() <= ty);
         Self::from_i256(self.as_i256(), ty)
     }
 
     pub fn zext(self, ty: Type) -> Self {
-        debug_assert!(self.ty() < ty);
+        debug_assert!(self.ty() <= ty);
         let i256: I256 = match self {
             Self::I1(val) => (val as u8).into(),
             Self::I8(val) => (val as u8).into(),

crates/parser/src/inst/data.rs

@@ -5,7 +5,7 @@ use crate::{ast, error::ArityBound, BuildCtx, Error};
 
 super::impl_inst_build! {Mload, has_mload, (addr: ValueId, ty: Type)}
 super::impl_inst_build! {Mstore, has_mstore, (value: ValueId, addr: ValueId, ty: Type)}
-super::impl_inst_build_common! {Gep, has_gep, ArityBound::AtLeast(0), build_gep}
+super::impl_inst_build_common! {Gep, has_gep, ArityBound::AtLeast(2), build_gep}
 
 fn build_gep(
     ctx: &mut BuildCtx,