escape analysis implementation

buggy for now

Cargo.lock

@@ -524,6 +524,7 @@ dependencies = [
  "parking_lot 0.11.2",
  "rayon",
  "regex",
+ "slab",
  "test-case",
 ]
 
@@ -1003,6 +1004,12 @@ version = "0.3.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7bd3e3206899af3f8b12af284fafc038cc1dc2b41d1b89dd17297221c5d225de"
 
+[[package]]
+name = "slab"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eb703cfe953bccee95685111adeedb76fabe4e97549a58d16f03ea7b9367bb32"
+
 [[package]]
 name = "smallvec"
 version = "1.8.0"

nac3core/Cargo.toml

@@ -9,6 +9,7 @@ itertools = "0.10.1"
 crossbeam = "0.8.1"
 parking_lot = "0.11.1"
 rayon = "1.5.1"
+slab = "0.4.6"
 nac3parser = { path = "../nac3parser" }
 
 [dependencies.inkwell]

nac3core/src/toplevel/composer.rs

@@ -4,7 +4,7 @@ use std::rc::Rc;
 use crate::{
     codegen::{expr::get_subst_key, stmt::exn_constructor},
     symbol_resolver::SymbolValue,
-    typecheck::type_inferencer::{FunctionData, Inferencer},
+    typecheck::{type_inferencer::{FunctionData, Inferencer}, escape_analysis::EscapeAnalyzer},
 };
 
 use super::*;
@@ -1792,6 +1792,7 @@ impl TopLevelComposer {
                             result
                         };
                         let mut calls: HashMap<CodeLocation, CallId> = HashMap::new();
+                        let mut args = vec![];
                         let mut inferencer = Inferencer {
                             top_level: ctx.as_ref(),
                             defined_identifiers: identifiers.clone(),
@@ -1812,6 +1813,7 @@ impl TopLevelComposer {
                                     result.insert("self".into(), self_ty);
                                 }
                                 result.extend(inst_args.iter().map(|x| (x.name, x.ty)));
+                                args.extend(result.iter().map(|(&a, &b)| (a, b)));
                                 result
                             },
                             primitives: primitives_ty,
@@ -1917,6 +1919,18 @@ impl TopLevelComposer {
                             ));
                         }
 
+                        if simple_name.to_string() != "__init__" {
+                            EscapeAnalyzer::check_function_lifetime(
+                                unifier,
+                                &primitives_ty,
+                                resolver.as_ref().unwrap().clone(),
+                                ctx.as_ref(),
+                                &args,
+                                &fun_body,
+                                ast.as_ref().unwrap().location,
+                            ).map_err(|e| format!("Escape analysis error: {}\n  in function {}", e, name))?;
+                        }
+
                         instance_to_stmt.insert(
                             get_subst_key(unifier, self_type, &subst, Some(&vars.keys().cloned().collect())),
                             FunInstance {

nac3core/src/typecheck/escape_analysis/lifetime.rs

@@ -1,276 +1,496 @@
-use std::cell::RefCell;
+use slab::Slab;
+use std::borrow::Cow;
 use std::collections::{HashMap, HashSet};
-use std::rc::Rc;
 
-use crate::typecheck::unification_table::{UnificationKey, UnificationTable};
+use nac3parser::ast::{Location, StrRef};
 
-use itertools::Itertools;
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-use nac3parser::ast::StrRef;
-
-// change this to enum, only local needs unification key
-pub type Lifetime = UnificationKey;
-
-#[derive(Copy, Debug, Clone, PartialEq, Eq, Hash)]
 pub enum LifetimeKind {
-    // can be assigned to fields of anything
+    Static,
-    // can be returned
-    // lifetime of static values
-    Global,
-    // can only be assigned to fields of objects with local lifetime
-    // can be returned
-    // lifetime of parameters
     NonLocal,
-    // can only be assigned to fields of objects with local lifetime
-    // cannot be returned
-    // lifetime of constructor return values
-    Local,
-    // can only be assigned to fields of objects with local lifetime
-    // cannot be returned
-    // lifetime of function return values
     Unknown,
+    PreciseLocal,
+    ImpreciseLocal,
 }
 
 impl std::ops::BitAnd for LifetimeKind {
     type Output = Self;
 
-    fn bitand(self, other: Self) -> Self {
+    fn bitand(self, rhs: Self) -> Self::Output {
         use LifetimeKind::*;
-        match (self, other) {
+        match (self, rhs) {
             (x, y) if x == y => x,
-            (Global, NonLocal) | (NonLocal, Global) => NonLocal,
+            (PreciseLocal, ImpreciseLocal) | (ImpreciseLocal, PreciseLocal) => ImpreciseLocal,
+            (Static, NonLocal) | (NonLocal, Static) => NonLocal,
             _ => Unknown,
         }
     }
 }
 
-impl std::cmp::PartialOrd for LifetimeKind {
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+pub struct LifetimeId(usize);
-        use LifetimeKind::*;
+
-        match (*self, *other) {
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-            (x, y) if x == y => Some(std::cmp::Ordering::Equal),
+pub struct BasicBlockId(usize);
-            (Local, _) | (_, Global) => Some(std::cmp::Ordering::Less),
+
-            (_, Local) | (Global, _) => Some(std::cmp::Ordering::Greater),
+#[derive(Debug, Clone)]
-            _ => None,
+pub enum LifetimeIR {
-        }
+    VarAssign { var: StrRef, lifetime: LifetimeId },
-    }
+    VarAccess { var: StrRef },
+    FieldAssign { obj: LifetimeId, field: StrRef, new: LifetimeId },
+    FieldAccess { obj: LifetimeId, field: StrRef },
+    CreateLifetime { kind: LifetimeKind },
+    PassedToFunc { param_lifetimes: Vec<LifetimeId> },
+    UnifyLifetimes { lifetimes: Vec<LifetimeId> },
+    Branch { targets: Vec<BasicBlockId> },
+    Return { val: Option<LifetimeId> },
 }
 
-pub struct BlockLifetimeContext {
+pub struct LifetimeIRBuilder {
-    mapping: Vec<(Option<Lifetime>, Lifetime)>,
+    irs: Vec<Option<(LifetimeIR, Location)>>,
+    basic_blocks: Vec<Vec<usize>>,
+    current_block: BasicBlockId,
 }
 
-impl BlockLifetimeContext {
+impl LifetimeIRBuilder {
     pub fn new() -> Self {
-        BlockLifetimeContext { mapping: Vec::new() }
+        LifetimeIRBuilder {
-    }
+            irs: vec![None],
-
+            basic_blocks: vec![vec![]],
-    pub fn add_fresh(&mut self, lifetime: Lifetime) {
+            current_block: BasicBlockId(0),
-        self.mapping.push((None, lifetime));
         }
     }
 
-struct LifetimeEntry {
+    pub fn print_ir(&self) -> String {
+        let mut lines = vec![];
+        for (i, bb) in self.basic_blocks.iter().enumerate() {
+            if bb.is_empty() {
+                continue;
+            }
+            lines.push(format!("{}:", i));
+            for ir in bb.iter() {
+                if let Some((inst, loc)) = &self.irs[*ir] {
+                    lines.push(format!("  {}: {:?} ({})", *ir, inst, loc));
+                }
+            }
+        }
+        lines.join("\n")
+    }
+
+    pub fn append_ir(&mut self, inst: LifetimeIR, loc: Location) -> LifetimeId {
+        let id = self.irs.len();
+        self.irs.push(Some((inst, loc)));
+        self.basic_blocks[self.current_block.0].push(id);
+        LifetimeId(id)
+    }
+
+    pub fn append_block(&mut self) -> BasicBlockId {
+        let id = self.basic_blocks.len();
+        self.basic_blocks.push(vec![]);
+        BasicBlockId(id)
+    }
+
+    pub fn get_current_block(&self) -> BasicBlockId {
+        self.current_block
+    }
+
+    pub fn position_at_end(&mut self, id: BasicBlockId) {
+        self.current_block = id;
+    }
+
+    pub fn is_terminated(&self, id: BasicBlockId) -> bool {
+        let bb = &self.basic_blocks[id.0];
+        if bb.is_empty() {
+            false
+        } else {
+            matches!(
+                self.irs[*bb.last().unwrap()],
+                Some((LifetimeIR::Return { .. }, _)) | Some((LifetimeIR::Branch { .. }, _))
+            )
+        }
+    }
+
+    pub fn remove_empty_bb(&mut self) {
+        let mut destination_mapping = HashMap::new();
+        let basic_blocks = &mut self.basic_blocks;
+        let irs = &mut self.irs;
+        for (i, bb) in basic_blocks.iter_mut().enumerate() {
+            bb.retain(|&id| irs[id].is_some());
+            if bb.len() == 1 {
+                let id = bb.pop().unwrap();
+                let ir = irs[id].take().unwrap();
+                match ir.0 {
+                    LifetimeIR::Branch { targets } => {
+                        destination_mapping.insert(i, targets);
+                    }
+                    _ => unreachable!(),
+                }
+            }
+        }
+        let mut buffer = HashSet::new();
+        for bb in basic_blocks.iter_mut() {
+            if bb.is_empty() {
+                continue;
+            }
+            if let LifetimeIR::Branch { targets } =
+                &mut irs[*bb.last().unwrap()].as_mut().unwrap().0
+            {
+                buffer.clear();
+                let mut updated = false;
+                for target in targets.iter() {
+                    if let Some(dest) = destination_mapping.get(&target.0) {
+                        buffer.extend(dest.iter().cloned());
+                        updated = true;
+                    } else {
+                        buffer.insert(*target);
+                    }
+                }
+                if updated {
+                    targets.clear();
+                    targets.extend(buffer.iter().cloned());
+                }
+            }
+        }
+    }
+
+    pub fn analyze(&self) -> Result<(), String> {
+        let mut analyzers = HashMap::new();
+        analyzers.insert(0, LifetimeAnalyzer::new());
+        let mut worklist = vec![0];
+        let mut counter = 0;
+        while let Some(bb) = worklist.pop() {
+            counter += 1;
+            if counter > 100 {
+                return Err(format!("escape analyzer stuck in infinite loop?"));
+            }
+            let mut analyzer = analyzers.get(&bb).unwrap().clone();
+            let block = &self.basic_blocks[bb];
+            let ir_iter = block.iter().filter_map(|&id| {
+                self.irs[id].as_ref().map(|(ir, loc)| (LifetimeId(id), ir, *loc))
+            });
+            if let Some(branch) = analyzer.analyze_basic_block(ir_iter)? {
+                for &target in branch.iter() {
+                    if let Some(successor) = analyzers.get_mut(&target.0) {
+                        if successor.merge(&analyzer) {
+                            // changed
+                            worklist.push(target.0);
+                        }
+                    } else {
+                        analyzers.insert(target.0, analyzer.clone());
+                        worklist.push(target.0);
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct LifetimeStore {
     kind: LifetimeKind,
-    fields: RefCell<HashMap<StrRef, Lifetime>>,
+    fields: HashMap<StrRef, LifetimeId>,
+    lifetimes: HashSet<LifetimeId>,
 }
 
-pub struct LifetimeTable {
+#[derive(Debug, Clone)]
-    table: UnificationTable<Rc<LifetimeEntry>>,
+pub struct LifetimeAnalyzer<'a> {
-    cache: HashSet<(Lifetime, Lifetime)>,
+    lifetime_to_id: HashMap<LifetimeId, usize>,
+    lifetime_stores: Slab<Cow<'a, LifetimeStore>>,
+    variable_assignment: HashMap<StrRef, LifetimeId>,
 }
 
-impl LifetimeTable {
+impl<'a> LifetimeAnalyzer<'a> {
     pub fn new() -> Self {
-        let mut zelf = Self { table: UnificationTable::new(), cache: Default::default() };
+        let mut zelf = LifetimeAnalyzer {
-        zelf.table.new_key(Rc::new(LifetimeEntry {
+            lifetime_to_id: HashMap::new(),
-            kind: LifetimeKind::Unknown,
+            lifetime_stores: Default::default(),
-            fields: Default::default(),
+            variable_assignment: HashMap::new(),
-        }));
+        };
+        zelf.add_lifetime(LifetimeId(0), LifetimeKind::Unknown);
         zelf
     }
 
-    pub fn add_lifetime(&mut self, kind: LifetimeKind) -> Lifetime {
+    pub fn merge(&mut self, other: &LifetimeAnalyzer) -> bool {
-        self.table.new_key(Rc::new(LifetimeEntry { kind, fields: Default::default() }))
+        let mut to_be_merged = other.lifetime_to_id.keys().cloned().collect::<Vec<_>>();
-    }
+        let mut updated = false;
 
-    pub fn unify(&mut self, a: Lifetime, b: Lifetime, ctx: &mut BlockLifetimeContext) {
+        let mut lifetime_merge_list = vec![];
-        self.cache.clear();
+        for (&var_name, &lifetime) in other.variable_assignment.iter() {
-        self.unify_impl(a, b, ctx);
+            if let Some(&our_lifetime) = self.variable_assignment.get(&var_name) {
+                if our_lifetime != lifetime {
+                    lifetime_merge_list.push((our_lifetime, lifetime));
                 }
-
-    fn get_scoped<const N: usize>(
-        &mut self,
-        mut lifetimes: [Lifetime; N],
-        ctx: &mut BlockLifetimeContext,
-    ) -> [Lifetime; N] {
-        for l in lifetimes.iter_mut() {
-            let mut result = None;
-            for (k, v) in ctx.mapping.iter() {
-                if self.table.unioned(*v, *l) || k.map_or(false, |k| self.table.unioned(k, *l)) {
-                    // already fresh
-                    result = Some(*v);
-                    break;
-                }
-            }
-            if let Some(result) = result {
-                *l = result;
             } else {
-                let lifetime = self.table.probe_value(*l).clone();
+                self.variable_assignment.insert(var_name, lifetime);
-                *l = if lifetime.kind == LifetimeKind::Unknown {
+                updated = true;
-                    UnificationKey(0)
+            }
+        }
+
+        while let Some(lifetime) = to_be_merged.pop() {
+            let other_store_id = *other.lifetime_to_id.get(&lifetime).unwrap();
+            if let Some(&self_store_id) = self.lifetime_to_id.get(&lifetime) {
+                let self_store = self.lifetime_stores.get_mut(self_store_id).unwrap();
+                let other_store = other.lifetime_stores.get(other_store_id).unwrap();
+                let self_store = self_store.to_mut();
+                // merge them
+                for (&field, &other_lifetime) in other_store.fields.iter() {
+                    if let Some(&self_lifetime) = self_store.fields.get(&field) {
+                        if self_lifetime != other_lifetime {
+                            lifetime_merge_list.push((self_lifetime, other_lifetime));
+                        }
                     } else {
-                    let k = self.table.new_key(lifetime);
+                        self_store.fields.insert(field, other_lifetime);
-                    ctx.mapping.push((Some(*l), k));
+                        updated = true;
-                    k
                     }
                 }
+                let zelf_lifetimes = &mut self_store.lifetimes;
+                for &other_lifetime in other_store.lifetimes.iter() {
+                    if zelf_lifetimes.insert(other_lifetime) {
+                        lifetime_merge_list.push((lifetime, other_lifetime));
+                    }
+                }
+                let result_kind = self_store.kind & other_store.kind;
+                if self_store.kind != result_kind {
+                    self_store.kind = result_kind;
+                }
+            } else {
+                let store = other.lifetime_stores.get(other_store_id).unwrap().as_ref().clone();
+                let store = self.lifetime_stores.insert(Cow::Owned(store));
+                self.lifetime_to_id.insert(lifetime, store);
+                updated = true;
             }
-        lifetimes
         }
 
-    fn unify_impl(&mut self, a: Lifetime, b: Lifetime, ctx: &mut BlockLifetimeContext) {
+        for (a, b) in lifetime_merge_list.into_iter() {
-        use LifetimeKind::*;
+            self.unify(a, b);
+        }
 
-        let [a, b] = self.get_scoped([a, b], ctx);
+        updated
-        let a = self.table.get_representative(a);
+    }
-        let b = self.table.get_representative(b);
+
-        if a == b || self.cache.contains(&(a, b)) || self.cache.contains(&(b, a)) {
+    pub fn add_lifetime(&mut self, lifetime: LifetimeId, kind: LifetimeKind) {
+        let id = self.lifetime_stores.insert(Cow::Owned(LifetimeStore {
+            kind,
+            fields: HashMap::new(),
+            lifetimes: [lifetime].iter().cloned().collect(),
+        }));
+        let old_store_id = self.lifetime_to_id.insert(lifetime, id);
+        if let Some(old_store_id) = old_store_id {
+            let old_lifetime_store = self.lifetime_stores.get_mut(old_store_id).unwrap().to_mut();
+            old_lifetime_store.lifetimes.remove(&lifetime);
+            if old_lifetime_store.lifetimes.is_empty() {
+                self.lifetime_stores.remove(old_store_id);
+            }
+        }
+    }
+
+    pub fn set_lifetime(&mut self, lifetime: LifetimeId, to: LifetimeId) {
+        let id = *self.lifetime_to_id.get(&to).unwrap();
+        let store = self.lifetime_stores.get_mut(id).unwrap();
+        store.to_mut().lifetimes.insert(lifetime);
+        let old_store_id = self.lifetime_to_id.insert(lifetime, id);
+        if let Some(old_store_id) = old_store_id {
+            let old_lifetime_store = self.lifetime_stores.get_mut(old_store_id).unwrap().to_mut();
+            old_lifetime_store.lifetimes.remove(&lifetime);
+            if old_lifetime_store.lifetimes.is_empty() {
+                self.lifetime_stores.remove(old_store_id);
+            }
+        }
+    }
+
+    fn unify(&mut self, lhs: LifetimeId, rhs: LifetimeId) {
+        use LifetimeKind::{ImpreciseLocal, PreciseLocal};
+        let lhs_id = *self.lifetime_to_id.get(&lhs).unwrap();
+        let rhs_id = *self.lifetime_to_id.get(&rhs).unwrap();
+        if lhs_id == rhs_id {
             return;
         }
-        self.cache.insert((a, b));
+        let lhs_store = self.lifetime_stores.get(lhs_id).unwrap();
-
+        let rhs_store = self.lifetime_stores.get(rhs_id).unwrap();
-        let v_a = self.table.probe_value(a).clone();
+        let all_lifetimes: HashSet<_> =
-        let v_b = self.table.probe_value(b).clone();
+            lhs_store.lifetimes.union(&rhs_store.lifetimes).cloned().collect();
-
+        let result_kind = lhs_store.kind & rhs_store.kind;
-        let result_kind = v_a.kind & v_b.kind;
+        let fields = if matches!(result_kind, PreciseLocal | ImpreciseLocal) {
-
+            let mut need_union = vec![];
-        let fields = if result_kind == Local {
+            let mut fields = lhs_store.fields.clone();
-            // we only need to track fields lifetime for objects with local lifetime
+            for (k, v) in rhs_store.fields.iter() {
-            let fields = v_a.fields.clone();
+                if let Some(old) = fields.insert(*k, *v) {
-            {
+                    need_union.push((old, *v));
-                let mut fields_ref = fields.borrow_mut();
-                for (k, v) in v_b.fields.borrow().iter() {
-                    if let Some(old) = fields_ref.insert(k.clone(), *v) {
-                        self.unify_impl(old, *v, ctx);
                 }
             }
+            drop(lhs_store);
+            drop(rhs_store);
+            for (lhs, rhs) in need_union {
+                self.unify(lhs, rhs);
             }
             fields
         } else {
             Default::default()
         };
-
+        // unify them, slow
-        self.table.unify(a, b);
+        for lifetime in all_lifetimes.iter() {
-        self.table.set_value(a, Rc::new(LifetimeEntry { kind: result_kind, fields }));
+            self.lifetime_to_id.insert(*lifetime, lhs_id);
+        }
+        *self.lifetime_stores.get_mut(lhs_id).unwrap() =
+            Cow::Owned(LifetimeStore { kind: result_kind, fields, lifetimes: all_lifetimes });
+        self.lifetime_stores.remove(rhs_id);
     }
 
-    pub fn get_field_lifetime(
+    fn get_field_lifetime(&self, obj: LifetimeId, field: StrRef) -> LifetimeId {
-        &mut self,
-        lifetime: Lifetime,
-        field: StrRef,
-        ctx: &mut BlockLifetimeContext,
-    ) -> Lifetime {
         use LifetimeKind::*;
-        let [lifetime] = self.get_scoped([lifetime], ctx);
+        let id = *self.lifetime_to_id.get(&obj).unwrap();
-        if let LifetimeEntry { kind: Local, fields } = &*self.table.probe_value(lifetime).clone() {
+        let store = &self.lifetime_stores.get(id).unwrap();
-            if let Some(lifetime) = fields.borrow().get(&field) {
+        if let Some(lifetime) = store.fields.get(&field) {
             *lifetime
+        } else if matches!(store.kind, PreciseLocal | ImpreciseLocal) {
+            LifetimeId(0)
         } else {
-                // unknown lifetime
+            obj
-                // we can reuse this lifetime because it will never be unified to something else
-                UnificationKey(0)
-            }
-        } else {
-            lifetime
         }
     }
 
-    pub fn set_field_lifetime(
+    fn set_field_lifetime(
         &mut self,
-        obj: Lifetime,
+        obj: LifetimeId,
         field: StrRef,
-        lifetime: Lifetime,
+        field_lifetime: LifetimeId,
-        is_strong_update: bool,
-        ctx: &mut BlockLifetimeContext,
     ) -> Result<(), String> {
-        let [obj, lifetime] = self.get_scoped([obj, lifetime], ctx);
+        use LifetimeKind::*;
-        let obj_lifetime = self.table.probe_value(obj).clone();
+        let obj_id = *self.lifetime_to_id.get(&obj).unwrap();
-        let field_lifetime = self.table.probe_value(lifetime).clone();
+        let field_id = *self.lifetime_to_id.get(&field_lifetime).unwrap();
-        if !(obj_lifetime.kind <= field_lifetime.kind) {
+        let field_lifetime_kind = self.lifetime_stores.get(field_id).unwrap().kind;
-            return Err("lifetime error".to_string());
+        let obj_store = self.lifetime_stores.get_mut(obj_id).unwrap();
+        if !matches!(
+            (obj_store.kind, field_lifetime_kind),
+            (PreciseLocal, _) | (ImpreciseLocal, _) | (_, Static)
+        ) {
+            return Err("field lifetime error".into());
         }
-        let mut fields = obj_lifetime.fields.borrow_mut();
+        match obj_store.kind {
-        if is_strong_update {
+            // $elem means list elements
-            fields.insert(field, lifetime);
+            PreciseLocal if field != "$elem".into() => {
+                // strong update
+                obj_store.to_mut().fields.insert(field, field_lifetime);
+            }
+            PreciseLocal | ImpreciseLocal => {
+                // weak update
+                let old_lifetime = obj_store
+                    .to_mut()
+                    .fields
+                    .get(&field)
+                    .copied();
+                if let Some(old_lifetime) = old_lifetime {
+                    self.unify(old_lifetime, field_lifetime);
                 } else {
-            if let Some(old) = fields.insert(field, lifetime) {
+                    obj_store.to_mut().fields.insert(field, field_lifetime);
-                self.unify(old, lifetime, ctx);
                 }
             }
+            _ => (),
+        }
         Ok(())
     }
 
-    pub fn get_lifetime_kind(
+    fn get_lifetime_kind(&self, lifetime: LifetimeId) -> LifetimeKind {
-        &mut self,
+        self.lifetime_stores.get(*self.lifetime_to_id.get(&lifetime).unwrap()).unwrap().kind
-        lifetime: Lifetime,
-        ctx: &mut BlockLifetimeContext,
-    ) -> LifetimeKind {
-        let [lifetime] = self.get_scoped([lifetime], ctx);
-        self.table.probe_value(lifetime).kind
     }
 
-    pub fn set_function_params(&mut self, lifetime: Lifetime, ctx: &mut BlockLifetimeContext) {
+    fn pass_function_params(&mut self, lifetimes: &[LifetimeId]) {
         use LifetimeKind::*;
-        // unify each field with global
+        let mut visited = HashSet::new();
-        let [lifetime] = self.get_scoped([lifetime], ctx);
+        let mut worklist = vec![];
-        let lifetime = self.table.probe_value(lifetime).clone();
+
-        let mut worklist = lifetime.fields.borrow().values().copied().collect_vec();
+        fn add_fields_to_worklist(
-        while let Some(item) = worklist.pop() {
+            visited: &mut HashSet<LifetimeId>,
-            let [item] = self.get_scoped([item], ctx);
+            worklist: &mut Vec<(LifetimeId, bool)>,
-            let lifetime = self.table.probe_value(item).clone();
+            fields: &HashMap<StrRef, LifetimeId>,
-            if lifetime.kind == Unknown || lifetime.kind == Global {
+        ) {
+            for (&name, &field) in fields.iter() {
+                if visited.insert(field) {
+                    // not visited previously
+                    let name = name.to_string();
+                    let mutable = !(name.starts_with("$elem") && name.len() != "$elem".len());
+                    worklist.push((field, mutable));
+                }
+            }
+        }
+
+        for lifetime in lifetimes.iter() {
+            let lifetime =
+                self.lifetime_stores.get_mut(*self.lifetime_to_id.get(lifetime).unwrap()).unwrap();
+            add_fields_to_worklist(&mut visited, &mut worklist, &lifetime.fields);
+        }
+        while let Some((item, mutable)) = worklist.pop() {
+            let lifetime =
+                self.lifetime_stores.get_mut(*self.lifetime_to_id.get(&item).unwrap()).unwrap();
+            if matches!(lifetime.kind, Unknown | Static) {
                 continue;
             }
-            let fields = lifetime.fields.borrow().clone();
+            add_fields_to_worklist(&mut visited, &mut worklist, &lifetime.fields);
-            for (_, v) in fields.iter() {
+            if mutable {
-                worklist.push(*v);
+                // we may assign values with static lifetime to function params
+                lifetime.to_mut().kind = lifetime.kind & Static;
             }
-            self.table.set_value(
-                item,
-                Rc::new(LifetimeEntry {
-                    kind: lifetime.kind & Global,
-                    fields: RefCell::new(fields),
-                }),
-            );
         }
     }
 
-    pub fn get_unknown_lifetime(&self) -> Lifetime {
+    pub fn analyze_basic_block<'b, I: Iterator<Item = (LifetimeId, &'b LifetimeIR, Location)>>(
-        UnificationKey(0)
+        &mut self,
+        instructions: I,
+    ) -> Result<Option<&'b [BasicBlockId]>, String> {
+        use LifetimeIR::*;
+        for (id, inst, loc) in instructions {
+            match inst {
+                VarAssign { var, lifetime } => {
+                    self.variable_assignment.insert(*var, *lifetime);
                 }
-
+                VarAccess { var } => {
-    pub fn equiv(&mut self, a: Lifetime, b: Lifetime, ctx: &mut BlockLifetimeContext) -> bool {
+                    let lifetime = self.variable_assignment.get(var).cloned();
-        use LifetimeKind::Local;
+                    if let Some(lifetime) = lifetime {
-        let [a, b] = self.get_scoped([a, b], ctx);
+                        self.set_lifetime(id, lifetime);
-        if self.table.unioned(a, b) {
-            return true;
-        }
-        let lifetime_a = self.table.probe_value(a).clone();
-        let lifetime_b = self.table.probe_value(b).clone();
-        if lifetime_a.kind == Local && lifetime_b.kind == Local {
-            let fields_a = lifetime_a.fields.borrow();
-            let fields_b = lifetime_b.fields.borrow();
-            for (k, v) in fields_a.iter() {
-                if fields_b.get(k).map(|v1| self.equiv(*v, *v1, ctx)) != Some(true) {
-                    return false;
-                }
-            }
-            // they are just equivalent
-            // this can avoid infinite recursion
-            self.table.unify(a, b);
-            true
                     } else {
-            lifetime_a.kind == lifetime_b.kind
+                        // should be static lifetime
+                        self.add_lifetime(id, LifetimeKind::Static)
                     }
                 }
+                FieldAssign { obj, field, new } => {
+                    self.set_field_lifetime(*obj, *field, *new)
+                        .map_err(|e| format!("{} in {}", e, loc))?;
+                }
+                FieldAccess { obj, field } => {
+                    let lifetime = self.get_field_lifetime(*obj, *field);
+                    self.set_lifetime(id, lifetime);
+                }
+                CreateLifetime { kind } => {
+                    if *kind == LifetimeKind::Unknown {
+                        self.set_lifetime(id, LifetimeId(0));
+                    } else {
+                        self.add_lifetime(id, *kind);
+                    }
+                }
+                PassedToFunc { param_lifetimes } => {
+                    self.pass_function_params(param_lifetimes);
+                }
+                UnifyLifetimes { lifetimes } => {
+                    assert!(!lifetimes.is_empty());
+                    let lhs = lifetimes[0];
+                    for rhs in lifetimes[1..].iter() {
+                        self.unify(lhs, *rhs);
+                    }
+                    self.set_lifetime(id, lhs);
+                }
+                Return { val } => {
+                    if let Some(val) = val {
+                        let kind = self.get_lifetime_kind(*val);
+                        if !matches!(kind, LifetimeKind::Static | LifetimeKind::NonLocal) {
+                            return Err(format!("return value lifetime error in {}", loc));
+                        }
+                    }
+                    return Ok(None);
+                }
+                Branch { targets } => return Ok(Some(targets))
+            }
+        }
+        Ok(None)
+    }
 }

nac3core/src/typecheck/escape_analysis/mod.rs

@@ -1,45 +1,46 @@
-use std::{collections::HashMap, sync::Arc};
+use std::sync::Arc;
 
-use nac3parser::ast::{Constant, Expr, ExprKind, Stmt, StmtKind, StrRef};
+use itertools::chain;
+use nac3parser::ast::{Comprehension, Constant, Expr, ExprKind, Location, Stmt, StmtKind, StrRef};
+
+use lifetime::{BasicBlockId, LifetimeIR, LifetimeIRBuilder, LifetimeId, LifetimeKind};
 
 use crate::{
     symbol_resolver::SymbolResolver,
     toplevel::{TopLevelContext, TopLevelDef},
 };
 
-use self::lifetime::{BlockLifetimeContext, Lifetime, LifetimeTable};
-
 use super::{
     type_inferencer::PrimitiveStore,
     typedef::{Type, TypeEnum, Unifier},
 };
 
-pub mod lifetime;
-
 #[cfg(test)]
 mod test;
 
-struct LifetimeContext<'a> {
+mod lifetime;
-    variable_mapping: HashMap<StrRef, (Lifetime, bool)>,
+
-    scope_ctx: BlockLifetimeContext,
+pub struct EscapeAnalyzer<'a> {
-    lifetime_table: LifetimeTable,
+    builder: LifetimeIRBuilder,
-    primitive_store: &'a PrimitiveStore,
+    loop_head: Option<BasicBlockId>,
+    loop_tail: Option<BasicBlockId>,
     unifier: &'a mut Unifier,
+    primitive_store: &'a PrimitiveStore,
     resolver: Arc<dyn SymbolResolver + Send + Sync>,
     top_level: &'a TopLevelContext,
 }
 
-impl<'a> LifetimeContext<'a> {
+impl<'a> EscapeAnalyzer<'a> {
     pub fn new(
         unifier: &'a mut Unifier,
         primitive_store: &'a PrimitiveStore,
         resolver: Arc<dyn SymbolResolver + Send + Sync>,
         top_level: &'a TopLevelContext,
-    ) -> LifetimeContext<'a> {
+    ) -> Self {
-        LifetimeContext {
+        Self {
-            variable_mapping: HashMap::new(),
+            builder: LifetimeIRBuilder::new(),
-            scope_ctx: BlockLifetimeContext::new(),
+            loop_head: None,
-            lifetime_table: LifetimeTable::new(),
+            loop_tail: None,
             primitive_store,
             unifier,
             resolver,
@@ -47,272 +48,305 @@ impl<'a> LifetimeContext<'a> {
         }
     }
 
-    fn get_expr_lifetime(
+    pub fn check_function_lifetime(
-        &mut self,
+        unifier: &'a mut Unifier,
-        expr: &Expr<Option<Type>>,
+        primitive_store: &'a PrimitiveStore,
-    ) -> Result<Option<(Lifetime, bool)>, String> {
+        resolver: Arc<dyn SymbolResolver + Send + Sync>,
-        let ty = expr.custom.unwrap();
+        top_level: &'a TopLevelContext,
-        let is_primitive = self.unifier.unioned(ty, self.primitive_store.int32)
+        args: &[(StrRef, Type)],
+        body: &[Stmt<Option<Type>>],
+        loc: Location,
+    ) -> Result<(), String> {
+        use LifetimeIR::{CreateLifetime, VarAssign};
+        let mut zelf = Self::new(unifier, primitive_store, resolver, top_level);
+        let nonlocal_lifetime =
+            zelf.builder.append_ir(CreateLifetime { kind: LifetimeKind::NonLocal }, loc);
+        for (name, ty) in args.iter().copied() {
+            if zelf.need_alloca(ty) {
+                zelf.builder.append_ir(VarAssign { var: name, lifetime: nonlocal_lifetime }, loc);
+            }
+        }
+        zelf.handle_statements(body)?;
+        zelf.builder.analyze()
+    }
+
+    fn need_alloca(&mut self, ty: Type) -> bool {
+        !(self.unifier.unioned(ty, self.primitive_store.int32)
             || self.unifier.unioned(ty, self.primitive_store.int64)
             || self.unifier.unioned(ty, self.primitive_store.uint32)
             || self.unifier.unioned(ty, self.primitive_store.uint64)
             || self.unifier.unioned(ty, self.primitive_store.float)
             || self.unifier.unioned(ty, self.primitive_store.bool)
             || self.unifier.unioned(ty, self.primitive_store.none)
-            || self.unifier.unioned(ty, self.primitive_store.range);
+            || self.unifier.unioned(ty, self.primitive_store.range))
+    }
 
+    fn is_terminated(&self) -> bool {
+        self.builder.is_terminated(self.builder.get_current_block())
+    }
+
+    fn handle_unknown_function_call<P: std::borrow::Borrow<Expr<Option<Type>>>>(
+        &mut self,
+        params: &[P],
+        ret_need_alloca: bool,
+        loc: Location,
+    ) -> Result<Option<LifetimeId>, String> {
+        let param_lifetimes = params
+            .iter()
+            .filter_map(|p| self.handle_expr(p.borrow()).transpose())
+            .collect::<Result<Vec<_>, _>>()?;
+        self.builder.append_ir(LifetimeIR::PassedToFunc { param_lifetimes }, loc);
+        if ret_need_alloca {
+            Ok(Some(
+                self.builder
+                    .append_ir(LifetimeIR::CreateLifetime { kind: LifetimeKind::Unknown }, loc),
+            ))
+        } else {
+            Ok(None)
+        }
+    }
+
+    fn handle_expr(&mut self, expr: &Expr<Option<Type>>) -> Result<Option<LifetimeId>, String> {
+        use LifetimeIR::*;
+        use LifetimeKind::*;
+        let need_alloca = self.need_alloca(expr.custom.unwrap());
+        let loc = expr.location;
         Ok(match &expr.node {
             ExprKind::Name { id, .. } => {
-                if let Some(lifetime) = self.variable_mapping.get(id) {
+                if need_alloca {
-                    Some(*lifetime)
+                    Some(self.builder.append_ir(VarAccess { var: *id }, loc))
                 } else {
-                    if is_primitive {
                     None
-                    } else {
-                        let lifetime =
-                            self.lifetime_table.add_lifetime(lifetime::LifetimeKind::Global);
-                        self.variable_mapping.insert(id.clone(), (lifetime, false));
-                        Some((lifetime, false))
-                    }
                 }
             }
             ExprKind::Attribute { value, attr, .. } => {
-                if is_primitive {
+                if need_alloca {
-                    self.get_expr_lifetime(value)?;
+                    let val = self.handle_expr(value)?.unwrap();
-                    None
+                    Some(self.builder.append_ir(FieldAccess { obj: val, field: *attr }, loc))
                 } else {
-                    self.get_expr_lifetime(value)?.map(|lifetime| {
+                    self.handle_expr(value)?;
-                        (
+                    None
-                            self.lifetime_table.get_field_lifetime(
-                                lifetime.0,
-                                *attr,
-                                &mut self.scope_ctx,
-                            ),
-                            false, // not sure if it is strong update for now...
-                        )
-                    })
                 }
             }
             ExprKind::Constant { .. } => {
-                if is_primitive {
+                if need_alloca {
-                    None
+                    Some(self.builder.append_ir(CreateLifetime { kind: Static }, loc))
                 } else {
-                    Some((self.lifetime_table.add_lifetime(lifetime::LifetimeKind::Global), false))
+                    None
                 }
             }
             ExprKind::List { elts, .. } => {
                 let elems =
-                    elts.iter()
+                    elts.iter().map(|e| self.handle_expr(e)).collect::<Result<Vec<_>, _>>()?;
-                        .map(|expr| self.get_expr_lifetime(expr))
+                let list_lifetime =
-                        .collect::<Result<Vec<_>, _>>()?;
+                    self.builder.append_ir(CreateLifetime { kind: PreciseLocal }, loc);
-                let elem = elems.into_iter().reduce(|prev, next| {
+                if !elems.is_empty() {
-                    if prev.is_some() {
+                    if elems[0].is_some() {
-                        self.lifetime_table.unify(
+                        let elems = elems.into_iter().map(|e| e.unwrap()).collect::<Vec<_>>();
-                            prev.unwrap().0,
+                        let elem_lifetime =
-                            next.unwrap().0,
+                            self.builder.append_ir(UnifyLifetimes { lifetimes: elems }, loc);
-                            &mut self.scope_ctx,
+                        self.builder.append_ir(
+                            FieldAssign {
+                                obj: list_lifetime,
+                                field: "$elem".into(),
+                                new: elem_lifetime,
+                            },
+                            loc,
                         );
                     }
-                    prev
-                });
-                let list_lifetime = self.lifetime_table.add_lifetime(lifetime::LifetimeKind::Local);
-
-                if let Some(Some(elem)) = elem {
-                    self.lifetime_table
-                        .set_field_lifetime(
-                            list_lifetime,
-                            "elem".into(),
-                            elem.0,
-                            true,
-                            &mut self.scope_ctx,
-                        )
-                        .unwrap();
-                }
-                Some((list_lifetime, true))
-            }
-            ExprKind::Subscript { value, slice, .. } => {
-                // value must be a list, so lifetime cannot be None
-                let (value_lifetime, _) = self.get_expr_lifetime(value)?.unwrap();
-                match &slice.node {
-                    ExprKind::Slice { lower, upper, step } => {
-                        for expr in [lower, upper, step].iter().filter_map(|x| x.as_ref()) {
-                            // account for side effects when computing the slice
-                            self.get_expr_lifetime(expr)?;
-                        }
-                        Some((
-                            self.lifetime_table.add_lifetime(lifetime::LifetimeKind::Local),
-                            true,
-                        ))
-                    }
-                    ExprKind::Constant { value: Constant::Int(v), .. } => {
-                        if is_primitive {
-                            None
-                        } else if let TypeEnum::TList { .. } =
-                            &*self.unifier.get_ty(value.custom.unwrap())
-                        {
-                            Some((
-                                self.lifetime_table.get_field_lifetime(
-                                    value_lifetime,
-                                    "elem".into(),
-                                    &mut self.scope_ctx,
-                                ),
-                                false,
-                            ))
                 } else {
-                            // tuple
+                    let elem_lifetime =
-                            Some((
+                        self.builder.append_ir(CreateLifetime { kind: PreciseLocal }, loc);
-                                self.lifetime_table.get_field_lifetime(
+                    self.builder.append_ir(
-                                    value_lifetime,
+                        FieldAssign {
-                                    format!("elem{}", v).into(),
+                            obj: list_lifetime,
-                                    &mut self.scope_ctx,
+                            field: "$elem".into(),
-                                ),
+                            new: elem_lifetime,
-                                false,
+                        },
-                            ))
+                        loc,
-                        }
+                    );
-                    }
-                    _ => {
-                        // account for side effects when computing the index
-                        self.get_expr_lifetime(slice)?;
-                        if is_primitive {
-                            None
-                        } else {
-                            Some((
-                                self.lifetime_table.get_field_lifetime(
-                                    value_lifetime,
-                                    "elem".into(),
-                                    &mut self.scope_ctx,
-                                ),
-                                false,
-                            ))
-                        }
-                    }
                 }
+                Some(list_lifetime)
             }
             ExprKind::Tuple { elts, .. } => {
                 let elems =
-                    elts.iter()
+                    elts.iter().map(|e| self.handle_expr(e)).collect::<Result<Vec<_>, _>>()?;
-                        .map(|expr| self.get_expr_lifetime(expr))
-                        .collect::<Result<Vec<_>, _>>()?;
                 let tuple_lifetime =
-                    self.lifetime_table.add_lifetime(lifetime::LifetimeKind::Local);
+                    self.builder.append_ir(CreateLifetime { kind: PreciseLocal }, loc);
                 for (i, lifetime) in elems.into_iter().enumerate() {
-                    if let Some((lifetime, _)) = lifetime {
+                    if let Some(lifetime) = lifetime {
-                        self.lifetime_table
+                        self.builder.append_ir(
-                            .set_field_lifetime(
+                            FieldAssign {
-                                tuple_lifetime,
+                                obj: tuple_lifetime,
-                                format!("elem{}", i).into(),
+                                field: format!("$elem{}", i).into(),
-                                lifetime,
+                                new: lifetime,
-                                true,
+                            },
-                                &mut self.scope_ctx,
+                            loc,
-                            )
+                        );
-                            .unwrap();
+                    }
+                }
+                Some(tuple_lifetime)
+            }
+            ExprKind::Subscript { value, slice, .. } => {
+                let value_lifetime = self.handle_expr(value)?.unwrap();
+                match &slice.node {
+                    ExprKind::Slice { lower, upper, step } => {
+                        for expr in [lower, upper, step].iter().filter_map(|x| x.as_ref()) {
+                            self.handle_expr(expr)?;
+                        }
+                        let slice_lifetime =
+                            self.builder.append_ir(CreateLifetime { kind: PreciseLocal }, loc);
+                        let slice_elem = self.builder.append_ir(
+                            FieldAccess { obj: value_lifetime, field: "$elem".into() },
+                            loc,
+                        );
+                        self.builder.append_ir(
+                            FieldAssign {
+                                obj: slice_lifetime,
+                                field: "$elem".into(),
+                                new: slice_elem,
+                            },
+                            loc,
+                        );
+                        Some(slice_lifetime)
+                    }
+                    ExprKind::Constant { value: Constant::Int(v), .. }
+                        if matches!(
+                            &*self.unifier.get_ty(value.custom.unwrap()),
+                            TypeEnum::TTuple { .. }
+                        ) =>
+                    {
+                        Some(self.builder.append_ir(
+                            FieldAccess {
+                                obj: value_lifetime,
+                                field: format!("$elem{}", v).into(),
+                            },
+                            loc,
+                        ))
+                    }
+                    _ => {
+                        self.handle_expr(slice)?;
+                        if need_alloca {
+                            Some(self.builder.append_ir(
+                                FieldAccess { obj: value_lifetime, field: "$elem".into() },
+                                loc,
+                            ))
+                        } else {
+                            None
+                        }
                     }
                 }
-                Some((tuple_lifetime, true))
             }
             ExprKind::Call { func, args, keywords } => {
-                let mut lifetimes = Vec::new();
+                let mut lifetimes = vec![];
-                for arg in args.iter() {
+                for arg in chain!(args.iter(), keywords.iter().map(|k| k.node.value.as_ref())) {
-                    if let Some(lifetime) = self.get_expr_lifetime(arg)? {
+                    if let Some(lifetime) = self.handle_expr(arg)? {
-                        lifetimes.push(lifetime.0);
+                        lifetimes.push(lifetime);
-                    }
-                }
-                for keyword in keywords.iter() {
-                    if let Some(lifetime) = self.get_expr_lifetime(&keyword.node.value)? {
-                        lifetimes.push(lifetime.0);
                     }
                 }
                 match &func.node {
                     ExprKind::Name { id, .. } => {
-                        for lifetime in lifetimes.into_iter() {
+                        self.builder.append_ir(PassedToFunc { param_lifetimes: lifetimes }, loc);
-                            self.lifetime_table.set_function_params(lifetime, &mut self.scope_ctx);
+                        if need_alloca {
-                        }
-                        if is_primitive {
-                            None
-                        } else {
                             let id = self
                                 .resolver
                                 .get_identifier_def(*id)
                                 .map_err(|e| format!("{} (at {})", e, func.location))?;
-                            // constructors
                             if let TopLevelDef::Class { .. } =
                                 &*self.top_level.definitions.read()[id.0].read()
                             {
-                                Some((
+                                Some(
-                                    self.lifetime_table.add_lifetime(lifetime::LifetimeKind::Local),
+                                    self.builder
-                                    true,
+                                        .append_ir(CreateLifetime { kind: PreciseLocal }, loc),
-                                ))
+                                )
                             } else {
-                                Some((self.lifetime_table.get_unknown_lifetime(), false))
+                                Some(self.builder.append_ir(CreateLifetime { kind: Unknown }, loc))
                             }
+                        } else {
+                            None
                         }
                     }
                     ExprKind::Attribute { value, .. } => {
-                        if let Some(lifetime) = self.get_expr_lifetime(value)? {
+                        let obj_lifetime = self.handle_expr(value)?.unwrap();
-                            lifetimes.push(lifetime.0);
+                        lifetimes.push(obj_lifetime);
-                        }
+                        self.builder.append_ir(PassedToFunc { param_lifetimes: lifetimes }, loc);
-                        for lifetime in lifetimes.into_iter() {
+                        if need_alloca {
-                            self.lifetime_table.set_function_params(lifetime, &mut self.scope_ctx);
+                            Some(self.builder.append_ir(CreateLifetime { kind: Unknown }, loc))
-                        }
-                        if is_primitive {
-                            None
                         } else {
-                            Some((self.lifetime_table.get_unknown_lifetime(), false))
+                            None
                         }
                     }
                     _ => unimplemented!(),
                 }
             }
-            ExprKind::BinOp { left, right, .. } => {
+            ExprKind::BinOp { left, right, .. } => self.handle_unknown_function_call(
-                let mut lifetimes = Vec::new();
+                &[left.as_ref(), right.as_ref()],
-                if let Some(l) = self.get_expr_lifetime(left)? {
+                need_alloca,
-                    lifetimes.push(l.0);
+                loc,
-                }
+            )?,
-                if let Some(l) = self.get_expr_lifetime(right)? {
-                    lifetimes.push(l.0);
-                }
-                for lifetime in lifetimes.into_iter() {
-                    self.lifetime_table.set_function_params(lifetime, &mut self.scope_ctx);
-                }
-                if is_primitive {
-                    None
-                } else {
-                    Some((self.lifetime_table.get_unknown_lifetime(), false))
-                }
-            }
             ExprKind::BoolOp { values, .. } => {
-                for v in values {
+                self.handle_unknown_function_call(&values, need_alloca, loc)?
-                    self.get_expr_lifetime(v)?;
-                }
-                None
             }
             ExprKind::UnaryOp { operand, .. } => {
-                if let Some(l) = self.get_expr_lifetime(operand)? {
+                self.handle_unknown_function_call(&[operand.as_ref()], need_alloca, loc)?
-                    self.lifetime_table.set_function_params(l.0, &mut self.scope_ctx);
-                }
-                if is_primitive {
-                    None
-                } else {
-                    Some((self.lifetime_table.get_unknown_lifetime(), false))
-                }
             }
             ExprKind::Compare { left, comparators, .. } => {
-                let mut lifetimes = Vec::new();
+                self.handle_unknown_function_call(&[left.as_ref()], false, loc)?;
-                if let Some(l) = self.get_expr_lifetime(left)? {
+                self.handle_unknown_function_call(&comparators, need_alloca, loc)?
-                    lifetimes.push(l.0);
             }
-                for c in comparators {
+            ExprKind::IfExp { test, body, orelse } => {
-                    if let Some(l) = self.get_expr_lifetime(c)? {
+                self.handle_expr(test)?;
-                        lifetimes.push(l.0);
+                let body_bb = self.builder.append_block();
-                    }
+                let else_bb = self.builder.append_block();
-                }
+                let tail_bb = self.builder.append_block();
-                for lifetime in lifetimes.into_iter() {
+                self.builder.append_ir(Branch { targets: vec![body_bb, else_bb] }, test.location);
-                    self.lifetime_table.set_function_params(lifetime, &mut self.scope_ctx);
+                self.builder.position_at_end(body_bb);
-                }
+                let body_lifetime = self.handle_expr(body)?;
-                // compare should give bool output, which does not have lifetime
+                self.builder.append_ir(Branch { targets: vec![tail_bb] }, body.location);
+                self.builder.position_at_end(else_bb);
+                let else_lifetime = self.handle_expr(body)?;
+                self.builder.append_ir(Branch { targets: vec![tail_bb] }, orelse.location);
+                self.builder.position_at_end(tail_bb);
+                if let (Some(body_lifetime), Some(else_lifetime)) = (body_lifetime, else_lifetime) {
+                    Some(self.builder.append_ir(
+                        UnifyLifetimes { lifetimes: vec![body_lifetime, else_lifetime] },
+                        loc,
+                    ))
+                } else {
                     None
                 }
-            // TODO: listcomp, ifexpr
+            }
+            ExprKind::ListComp { elt, generators } => {
+                let Comprehension { target, iter, ifs, .. } = &generators[0];
+                let list_lifetime =
+                    self.builder.append_ir(CreateLifetime { kind: PreciseLocal }, loc);
+                let iter_elem_lifetime = self.handle_expr(iter)?.map(|obj| {
+                    self.builder
+                        .append_ir(FieldAccess { obj, field: "$elem".into() }, iter.location)
+                });
+                let loop_body = self.builder.append_block();
+                let loop_tail = self.builder.append_block();
+                self.builder.append_ir(Branch { targets: vec![loop_body] }, loc);
+                self.builder.position_at_end(loop_body);
+                self.handle_assignment(target, iter_elem_lifetime)?;
+                for ifexpr in ifs.iter() {
+                    self.handle_expr(ifexpr)?;
+                }
+                let elem_lifetime = self.handle_expr(elt)?;
+                if let Some(elem_lifetime) = elem_lifetime {
+                    self.builder.append_ir(
+                        FieldAssign {
+                            obj: list_lifetime,
+                            field: "$elem".into(),
+                            new: elem_lifetime,
+                        },
+                        elt.location,
+                    );
+                }
+                self.builder.append_ir(Branch { targets: vec![loop_body, loop_tail] }, loc);
+                self.builder.position_at_end(loop_tail);
+                Some(list_lifetime)
+            }
             _ => unimplemented!(),
         })
     }
@@ -320,65 +354,63 @@ impl<'a> LifetimeContext<'a> {
     fn handle_assignment(
         &mut self,
         lhs: &Expr<Option<Type>>,
-        rhs_lifetime: Option<(Lifetime, bool)>,
+        rhs_lifetime: Option<LifetimeId>,
     ) -> Result<(), String> {
+        use LifetimeIR::*;
         match &lhs.node {
             ExprKind::Attribute { value, attr, .. } => {
-                let (lhs_lifetime, is_strong_update) = self.get_expr_lifetime(value)?.unwrap();
+                let value_lifetime = self.handle_expr(value)?.unwrap();
-                if let Some((lifetime, _)) = rhs_lifetime {
+                if let Some(field_lifetime) = rhs_lifetime {
-                    self.lifetime_table
+                    self.builder.append_ir(
-                        .set_field_lifetime(
+                        FieldAssign { obj: value_lifetime, field: *attr, new: field_lifetime },
-                            lhs_lifetime,
+                        lhs.location,
-                            *attr,
+                    );
-                            lifetime,
-                            is_strong_update,
-                            &mut self.scope_ctx,
-                        )
-                        .map_err(|_| format!("illegal field assignment in {}", lhs.location))?;
                 }
             }
             ExprKind::Subscript { value, slice, .. } => {
-                let (list_lifetime, _) = self.get_expr_lifetime(value)?.unwrap();
+                let value_lifetime = self.handle_expr(value)?.unwrap();
                 let elem_lifetime = if let ExprKind::Slice { lower, upper, step } = &slice.node {
-                    // compute side effects
                     for expr in [lower, upper, step].iter().filter_map(|x| x.as_ref()) {
-                        // account for side effects when computing the slice
+                        self.handle_expr(expr)?;
-                        self.get_expr_lifetime(expr)?;
                     }
-                    // slice assignment will copy elements from rhs to lhs
+                    if let Some(rhs_lifetime) = rhs_lifetime {
-                    self.lifetime_table.get_field_lifetime(
+                        // must be a list
-                        rhs_lifetime.unwrap().0,
+                        Some(self.builder.append_ir(
-                        "elem".into(),
+                            FieldAccess { obj: rhs_lifetime, field: "$elem".into() },
-                        &mut self.scope_ctx,
+                            lhs.location,
-                    )
+                        ))
                     } else {
-                    // must be list element, as assignment to tuple element is prohibited
+                        None
-                    self.get_expr_lifetime(slice)?;
+                    }
-                    rhs_lifetime.unwrap().0
+                } else {
+                    self.handle_expr(slice)?;
+                    rhs_lifetime
                 };
-                self.lifetime_table
+                // must be a list
-                    .set_field_lifetime(
+                if let Some(elem_lifetime) = elem_lifetime {
-                        list_lifetime,
+                    self.builder.append_ir(
-                        "elem".into(),
+                        FieldAssign {
-                        elem_lifetime,
+                            obj: value_lifetime,
-                        false,
+                            field: "$elem".into(),
-                        &mut self.scope_ctx,
+                            new: elem_lifetime,
-                    )
+                        },
-                    .map_err(|_| format!("illegal element assignment in {}", lhs.location))?;
+                        lhs.location,
+                    );
+                }
             }
             ExprKind::Name { id, .. } => {
                 if let Some(lifetime) = rhs_lifetime {
-                    self.variable_mapping.insert(*id, lifetime);
+                    self.builder.append_ir(VarAssign { var: *id, lifetime }, lhs.location);
                 }
             }
             ExprKind::Tuple { elts, .. } => {
+                let rhs_lifetime = rhs_lifetime.unwrap();
                 for (i, e) in elts.iter().enumerate() {
-                    let elem_lifetime = self.lifetime_table.get_field_lifetime(
+                    let elem_lifetime = self.builder.append_ir(
-                        rhs_lifetime.unwrap().0,
+                        FieldAccess { obj: rhs_lifetime, field: format!("$elem{}", i).into() },
-                        format!("elem{}", i).into(),
+                        e.location,
-                        &mut self.scope_ctx,
                     );
-                    self.handle_assignment(e, Some((elem_lifetime, false)))?;
+                    self.handle_assignment(e, Some(elem_lifetime))?;
                 }
             }
             _ => unreachable!(),
@@ -386,18 +418,151 @@ impl<'a> LifetimeContext<'a> {
         Ok(())
     }
 
-    pub fn handle_statement(&mut self, stmt: &Stmt<Option<Type>>) -> Result<(), String> {
+    fn handle_statement(&mut self, stmt: &Stmt<Option<Type>>) -> Result<(), String> {
+        use LifetimeIR::*;
         match &stmt.node {
             StmtKind::Expr { value, .. } => {
-                self.get_expr_lifetime(value)?;
+                self.handle_expr(value)?;
             }
             StmtKind::Assign { targets, value, .. } => {
-                let rhs_lifetime = self.get_expr_lifetime(value)?;
+                let rhs_lifetime = self.handle_expr(value)?;
-                for target in targets.iter() {
+                for target in targets {
                     self.handle_assignment(target, rhs_lifetime)?;
                 }
             }
-            _ => unimplemented!(),
+            StmtKind::If { test, body, orelse, .. } => {
+                // test should return bool
+                self.handle_expr(test)?;
+                let body_bb = self.builder.append_block();
+                let else_bb = self.builder.append_block();
+                self.builder.append_ir(Branch { targets: vec![body_bb, else_bb] }, stmt.location);
+                self.builder.position_at_end(body_bb);
+                self.handle_statements(&body)?;
+                let body_terminated = self.is_terminated();
+                if orelse.is_empty() {
+                    if !body_terminated {
+                        // else_bb is the basic block after this if statement
+                        self.builder.append_ir(Branch { targets: vec![else_bb] }, stmt.location);
+                        self.builder.position_at_end(else_bb);
+                    }
+                } else {
+                    let tail_bb = self.builder.append_block();
+                    if !body_terminated {
+                        self.builder.append_ir(Branch { targets: vec![tail_bb] }, stmt.location);
+                    }
+                    self.builder.position_at_end(else_bb);
+                    self.handle_statements(&orelse)?;
+                    if !self.is_terminated() {
+                        self.builder.append_ir(Branch { targets: vec![tail_bb] }, stmt.location);
+                    }
+                    self.builder.position_at_end(tail_bb);
+                }
+            }
+            StmtKind::While { test, body, orelse, .. } => {
+                let old_loop_head = self.loop_head;
+                let old_loop_tail = self.loop_tail;
+                let loop_head = self.builder.append_block();
+                let loop_body = self.builder.append_block();
+                let loop_else =
+                    if orelse.is_empty() { None } else { Some(self.builder.append_block()) };
+                let loop_tail = self.builder.append_block();
+                self.loop_head = Some(loop_head);
+                self.loop_tail = Some(loop_tail);
+                self.builder.append_ir(Branch { targets: vec![loop_head] }, stmt.location);
+                self.builder.position_at_end(loop_head);
+                self.handle_expr(test)?;
+                self.builder.append_ir(
+                    Branch { targets: vec![loop_body, loop_else.unwrap_or(loop_tail)] },
+                    stmt.location,
+                );
+                self.builder.position_at_end(loop_body);
+                self.handle_statements(&body)?;
+                if !self.is_terminated() {
+                    self.builder.append_ir(Branch { targets: vec![loop_head] }, stmt.location);
+                }
+
+                self.loop_head = old_loop_head;
+                self.loop_tail = old_loop_tail;
+                if let Some(loop_else) = loop_else {
+                    self.builder.position_at_end(loop_else);
+                    self.handle_statements(&orelse)?;
+                    if !self.is_terminated() {
+                        self.builder.append_ir(Branch { targets: vec![loop_tail] }, stmt.location);
+                    }
+                }
+                self.builder.position_at_end(loop_tail);
+            }
+            StmtKind::For { target, iter, body, orelse, .. } => {
+                let old_loop_head = self.loop_head;
+                let old_loop_tail = self.loop_tail;
+                let loop_head = self.builder.append_block();
+                let loop_body = self.builder.append_block();
+                let loop_else =
+                    if orelse.is_empty() { None } else { Some(self.builder.append_block()) };
+                let loop_tail = self.builder.append_block();
+                self.loop_head = Some(loop_head);
+                self.loop_tail = Some(loop_tail);
+                let iter_lifetime = self.handle_expr(iter)?.map(|obj| {
+                    self.builder
+                        .append_ir(FieldAccess { obj, field: "$elem".into() }, iter.location)
+                });
+                self.builder.append_ir(Branch { targets: vec![loop_head] }, stmt.location);
+                self.builder.position_at_end(loop_head);
+                if let Some(iter_lifetime) = iter_lifetime {
+                    self.handle_assignment(target, Some(iter_lifetime))?;
+                }
+                self.builder.append_ir(
+                    Branch { targets: vec![loop_body, loop_else.unwrap_or(loop_tail)] },
+                    stmt.location,
+                );
+                self.builder.position_at_end(loop_body);
+                self.handle_statements(&body)?;
+                if !self.is_terminated() {
+                    self.builder.append_ir(Branch { targets: vec![loop_head] }, stmt.location);
+                }
+
+                self.loop_head = old_loop_head;
+                self.loop_tail = old_loop_tail;
+                if let Some(loop_else) = loop_else {
+                    self.builder.position_at_end(loop_else);
+                    self.handle_statements(&orelse)?;
+                    if !self.is_terminated() {
+                        self.builder.append_ir(Branch { targets: vec![loop_tail] }, stmt.location);
+                    }
+                }
+                self.builder.position_at_end(loop_tail);
+            }
+
+            StmtKind::Continue { .. } => {
+                if let Some(loop_head) = self.loop_head {
+                    self.builder.append_ir(Branch { targets: vec![loop_head] }, stmt.location);
+                } else {
+                    return Err(format!("break outside loop"));
+                }
+            }
+            StmtKind::Break { .. } => {
+                if let Some(loop_tail) = self.loop_tail {
+                    self.builder.append_ir(Branch { targets: vec![loop_tail] }, stmt.location);
+                } else {
+                    return Err(format!("break outside loop"));
+                }
+            }
+            StmtKind::Return { value, .. } => {
+                let val = if let Some(value) = value { self.handle_expr(value)? } else { None };
+                self.builder.append_ir(Return { val }, stmt.location);
+            }
+            StmtKind::Pass { .. } => {}
+            _ => unimplemented!("{:?}", stmt.node),
+        }
+        Ok(())
+    }
+
+    fn handle_statements(&mut self, stmts: &[Stmt<Option<Type>>]) -> Result<(), String> {
+        for stmt in stmts.iter() {
+            if self.builder.is_terminated(self.builder.get_current_block()) {
+                break;
+            }
+            self.handle_statement(stmt)?;
         }
         Ok(())
     }

nac3core/src/typecheck/escape_analysis/test.rs

@@ -0,0 +1,60 @@
+use super::EscapeAnalyzer;
+use crate::typecheck::{type_inferencer::test::TestEnvironment, typedef::TypeEnum};
+use indoc::indoc;
+use nac3parser::ast::fold::Fold;
+use std::collections::hash_set::HashSet;
+use test_case::test_case;
+
+use nac3parser::parser::parse_program;
+
+#[test_case(indoc! {"
+    # a: list[list[int32]]
+    b = [1]
+    a[0] = b
+"}, Err("field lifetime error in unknown: line 3 column 2".into())
+; "assign global elem")]
+#[test_case(indoc! {"
+    # a: list[list[int32]]
+    b = [[], []]
+    b[1] = a
+    b[0][0] = [0]
+"}, Err("field lifetime error in unknown: line 4 column 5".into())
+; "global unify")]
+#[test_case(indoc! {"
+    b = [1, 2, 3]
+    c = [a]
+    c[0][0] = b
+"}, Err("field lifetime error in unknown: line 3 column 5".into())
+; "global unify 2")]
+fn test_simple(source: &str, expected_result: Result<(), String>) {
+    let mut env = TestEnvironment::basic_test_env();
+    let mut defined_identifiers: HashSet<_> = env.identifier_mapping.keys().cloned().collect();
+    defined_identifiers.insert("a".into());
+    let mut inferencer = env.get_inferencer();
+    inferencer.defined_identifiers = defined_identifiers.clone();
+
+    let list_int = inferencer.unifier.add_ty(TypeEnum::TList { ty: inferencer.primitives.int32 });
+    let list_list_int = inferencer.unifier.add_ty(TypeEnum::TList { ty: list_int });
+
+    inferencer.variable_mapping.insert("a".into(), list_list_int);
+    let statements = parse_program(source, Default::default()).unwrap();
+    let statements = statements
+        .into_iter()
+        .map(|v| inferencer.fold_stmt(v))
+        .collect::<Result<Vec<_>, _>>()
+        .unwrap();
+
+    inferencer.check_block(&statements, &mut defined_identifiers).unwrap();
+
+    let mut lifetime_ctx = EscapeAnalyzer::new(
+        &mut inferencer.unifier,
+        &mut inferencer.primitives,
+        inferencer.function_data.resolver.clone(),
+        &inferencer.top_level,
+    );
+    lifetime_ctx.handle_statements(&statements).unwrap();
+    lifetime_ctx.builder.remove_empty_bb();
+    let result = lifetime_ctx.builder.analyze();
+    assert_eq!(result, expected_result);
+
+}