From 722e3df0868be37c0740eb36346524494fd12b6a Mon Sep 17 00:00:00 2001
From: ychenfo <yc@m-labs.hk>
Date: Sun, 10 Apr 2022 16:25:15 +0800
Subject: [PATCH] nac3core, artiq: optimize kernel invariant for tuple index

---
 nac3artiq/src/symbol_resolver.rs | 22 ++++++++++++++++++++++
 nac3core/src/codegen/expr.rs     | 31 ++++++++++++++++++++++---------
 nac3core/src/symbol_resolver.rs  |  2 ++
 3 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/nac3artiq/src/symbol_resolver.rs b/nac3artiq/src/symbol_resolver.rs
index 416c81771..609de2c36 100644
--- a/nac3artiq/src/symbol_resolver.rs
+++ b/nac3artiq/src/symbol_resolver.rs
@@ -212,6 +212,28 @@ impl StaticValue for PythonValue {
             }))
         })
     }
+
+    fn get_tuple_element<'ctx>(&self, index: u32) -> Option<ValueEnum<'ctx>> {
+        Python::with_gil(|py| -> PyResult<Option<(u64, PyObject)>> {
+            let helper = &self.resolver.helper;
+            let ty = helper.type_fn.call1(py, (&self.value,))?;
+            let ty_id: u64 = helper.id_fn.call1(py, (ty,))?.extract(py)?;
+            assert_eq!(ty_id, self.resolver.primitive_ids.tuple);
+            let tup: &PyTuple = self.value.extract(py)?;
+            let elem = tup.get_item(index as usize);
+            let id = self.resolver.helper.id_fn.call1(py, (elem,))?.extract(py)?;
+            Ok(Some((id, elem.into())))
+        })
+        .unwrap()
+        .map(|(id, obj)| {
+            ValueEnum::Static(Arc::new(PythonValue {
+                id,
+                value: obj,
+                store_obj: self.store_obj.clone(),
+                resolver: self.resolver.clone(),
+            }))
+        })
+    }
 }
 
 impl InnerResolver {
diff --git a/nac3core/src/codegen/expr.rs b/nac3core/src/codegen/expr.rs
index 272249c21..3ca728355 100644
--- a/nac3core/src/codegen/expr.rs
+++ b/nac3core/src/codegen/expr.rs
@@ -1515,26 +1515,39 @@ pub fn gen_expr<'ctx, 'a, G: CodeGenerator>(
                         [Some(raw_index), Some(len), None],
                         expr.location,
                     );
-                    ctx.build_gep_and_load(arr_ptr, &[index])
+                    ctx.build_gep_and_load(arr_ptr, &[index]).into()
                 }
             } else if let TypeEnum::TTuple { .. } = &*ctx.unifier.get_ty(value.custom.unwrap()) {
-                let v = generator
-                    .gen_expr(ctx, value)?
-                    .unwrap()
-                    .to_basic_value_enum(ctx, generator, value.custom.unwrap())?
-                    .into_struct_value();
                 let index: u32 =
                     if let ExprKind::Constant { value: ast::Constant::Int(v), .. } = &slice.node {
                         (*v).try_into().unwrap()
                     } else {
                         unreachable!("tuple subscript must be const int after type check");
                     };
-                ctx.builder.build_extract_value(v, index, "tup_elem").unwrap()
+                let v = generator
+                    .gen_expr(ctx, value)?
+                    .unwrap();
+                match v {
+                    ValueEnum::Dynamic(v) => {
+                        let v = v.into_struct_value();
+                        ctx.builder.build_extract_value(v, index, "tup_elem").unwrap().into()
+                    }
+                    ValueEnum::Static(v) => {
+                        match v.get_tuple_element(index) {
+                            Some(v) => v,
+                            None => {
+                                let tup = v
+                                    .to_basic_value_enum(ctx, generator, value.custom.unwrap())?
+                                    .into_struct_value();
+                                ctx.builder.build_extract_value(tup, index, "tup_elem").unwrap().into()
+                            }
+                        }
+                    }
+                }
             } else {
                 unreachable!("should not be other subscriptable types after type check");
             }
-        }
-        .into(),
+        },
         ExprKind::ListComp { .. } => gen_comprehension(generator, ctx, expr)?.into(),
         _ => unimplemented!(),
     }))
diff --git a/nac3core/src/symbol_resolver.rs b/nac3core/src/symbol_resolver.rs
index acdd0d902..f6ce4fa3c 100644
--- a/nac3core/src/symbol_resolver.rs
+++ b/nac3core/src/symbol_resolver.rs
@@ -79,6 +79,8 @@ pub trait StaticValue {
         name: StrRef,
         ctx: &mut CodeGenContext<'ctx, 'a>,
     ) -> Option<ValueEnum<'ctx>>;
+
+    fn get_tuple_element<'ctx>(&self, index: u32) -> Option<ValueEnum<'ctx>>;
 }
 
 #[derive(Clone)]