From 259958aded5c9236ec926963ccd59c8e2b7e4d4e Mon Sep 17 00:00:00 2001
From: lyken <lyken@m-labs.hk>
Date: Sun, 14 Jul 2024 13:12:06 +0800
Subject: [PATCH] core: BoundedIxed Ixed & ArraySlice optics

---
 nac3core/src/codegen/optics.rs | 94 +++++++++++++++++++++++++++++++++-
 1 file changed, 93 insertions(+), 1 deletion(-)

diff --git a/nac3core/src/codegen/optics.rs b/nac3core/src/codegen/optics.rs
index e38d30aa..3ac7d954 100644
--- a/nac3core/src/codegen/optics.rs
+++ b/nac3core/src/codegen/optics.rs
@@ -6,7 +6,7 @@ use inkwell::{
 };
 use itertools::Itertools;
 
-use super::CodeGenContext;
+use super::{CodeGenContext, CodeGenerator};
 
 // TODO: Write a taxonomy
 
@@ -124,6 +124,12 @@ impl<'ctx, AddresseeOptic> OpticValue<'ctx> for Address<'ctx, AddresseeOptic> {
 #[derive(Debug, Clone)]
 pub struct AddressLens<AddresseeOptic>(pub AddresseeOptic);
 
+impl<AddresseeOptic> AddressLens<AddresseeOptic> {
+    pub fn new_opaque<'ctx>(&self, ctx: &CodeGenContext<'ctx, '_>) -> AddressLens<IntLens<'ctx>> {
+        AddressLens(IntLens(ctx.ctx.i8_type()))
+    }
+}
+
 impl<'ctx, AddresseeOptic: Optic<'ctx>> Optic<'ctx> for AddressLens<AddresseeOptic> {
     type Value = Address<'ctx, AddresseeOptic>;
 
@@ -341,3 +347,89 @@ impl<'ctx, AddresseeOptic: StructureOptic<'ctx>> Address<'ctx, AddresseeOptic> {
         field.get(ctx, self.address, field.name)
     }
 }
+
+// Name inspired by https://hackage.haskell.org/package/lens-5.3.2/docs/Control-Lens-At.html#t:Ixed
+pub trait Ixed<'ctx, ElementOptic> {
+    // TODO: Interface/Method to expose the IntType of index?
+    // or even make index itself parameterized? (probably no)
+
+    fn ix(
+        &self,
+        ctx: &CodeGenContext<'ctx, '_>,
+        index: IntValue<'ctx>,
+        name: &str,
+    ) -> Address<'ctx, ElementOptic>;
+}
+
+// TODO: Can do interface seggregation
+pub trait BoundedIxed<'ctx, ElementOptic>: Ixed<'ctx, ElementOptic> {
+    fn num_elements(&self, ctx: &CodeGenContext<'ctx, '_>) -> IntValue<'ctx>;
+
+    // Check if 0 <= index < self.num_elements()
+    fn ix_bounds_checked<G: CodeGenerator + ?Sized>(
+        &self,
+        generator: &mut G,
+        ctx: &mut CodeGenContext<'ctx, '_>,
+        index: IntValue<'ctx>,
+        name: &str,
+    ) -> Address<'ctx, ElementOptic> {
+        let num_elements = self.num_elements(ctx);
+        let int_type = num_elements.get_type(); // NOTE: Weird get_type(), see comment under `trait Ixed`
+
+        assert_eq!(int_type.get_bit_width(), index.get_type().get_bit_width()); // Might as well check bit width to catch bugs
+
+        // TODO: SGE or UGE? or make it defined by the implementee?
+
+        // Check `0 <= index`
+        let lower_bounded = ctx
+            .builder
+            .build_int_compare(
+                inkwell::IntPredicate::SLE,
+                int_type.const_zero(),
+                index,
+                "lower_bounded",
+            )
+            .unwrap();
+
+        // Check `index < num_elements`
+        let upper_bounded = ctx
+            .builder
+            .build_int_compare(inkwell::IntPredicate::SLT, index, num_elements, "upper_bounded")
+            .unwrap();
+
+        // Compute `0 <= index && index < num_elements`
+        let bounded = ctx.builder.build_and(lower_bounded, upper_bounded, "bounded").unwrap();
+
+        // Assert `bounded`
+        ctx.make_assert(generator, bounded, "0:IndexError", "nac3core LLVM codegen attempting to access out of bounds array index {0}. Must satisfy 0 <= index < {2}", [Some(index), Some(num_elements), None], ctx.current_loc);
+
+        // ...and finally do indexing
+        self.ix(ctx, index, name)
+    }
+}
+
+pub struct ArraySlice<'ctx, ElementOptic> {
+    pub num_elements: IntValue<'ctx>,
+    pub base: Address<'ctx, ElementOptic>,
+}
+
+impl<'ctx, ElementOptic: Optic<'ctx>> Ixed<'ctx, ElementOptic> for ArraySlice<'ctx, ElementOptic> {
+    fn ix(
+        &self,
+        ctx: &CodeGenContext<'ctx, '_>,
+        index: IntValue<'ctx>,
+        name: &str,
+    ) -> Address<'ctx, ElementOptic> {
+        let element_addr =
+            unsafe { ctx.builder.build_in_bounds_gep(self.base.address, &[index], name).unwrap() };
+        Address { address: element_addr, addressee_optic: self.base.addressee_optic.clone() }
+    }
+}
+
+impl<'ctx, ElementOptic: Optic<'ctx>> BoundedIxed<'ctx, ElementOptic>
+    for ArraySlice<'ctx, ElementOptic>
+{
+    fn num_elements(&self, _ctx: &CodeGenContext<'ctx, '_>) -> IntValue<'ctx> {
+        self.num_elements
+    }
+}