Record input types for loads and invokes in type propagator (#1339)

* Track input types for calls

* Simplify

* Allow block types

src/compiler/propagation/type_database.cc

@@ -45,9 +45,9 @@ TypeDatabase::~TypeDatabase() {
 
 void TypeDatabase::check_top(uint8* bcp, Object* value) const {
   int position = program_->absolute_bci_from_bcp(bcp);
-  auto probe = usage_.find(position);
-  if (probe == usage_.end()) {
-    FATAL("usage not analyzed: %d", position);
+  auto probe = output_.find(position);
+  if (probe == output_.end()) {
+    FATAL("output not analyzed: %d", position);
   }
   TypeSet type = probe->second;
   if (type.is_block()) {
@@ -159,9 +159,9 @@ const std::vector<TypeSet> TypeDatabase::arguments(Method method) const {
   return result;
 }
 
-const TypeSet TypeDatabase::usage(int position) const {
-  auto probe = usage_.find(position);
-  if (probe == usage_.end()) {
+const TypeSet TypeDatabase::output(int position) const {
+  auto probe = output_.find(position);
+  if (probe == output_.end()) {
     return TypeSet::invalid();
   } else {
     return probe->second;
@@ -209,15 +209,30 @@ std::string TypeDatabase::as_json() const {
   out << "[\n";
 
   bool first = true;
-  for (auto it : usage_) {
+  for (auto it : output_) {
     if (first) {
       first = false;
     } else {
       out << ",\n";
     }
-    std::string type_string = it.second.as_json(program_);
-    out << "  {\"position\": " << it.first;
-    out << ", \"type\": " << type_string << "}";
+
+    int position = it.first;
+    std::string output_string = it.second.as_json(program_);
+    out << "  {\"position\": " << position;
+    if (input_.find(position) != input_.end()) {
+      out << ", \"input\": [";
+      TypeStack* stack = input_.at(position);
+      for (int i = 0; i < stack->size(); i++) {
+        if (i != 0) {
+          out << ",";
+        }
+        TypeSet type = stack->get(i);
+        std::string type_string = type.as_json(program_);
+        out << type_string;
+      }
+      out << "]";
+    }
+    out << ", \"output\": " << output_string << "}";
   }
 
   for (auto it : methods_) {
@@ -263,11 +278,23 @@ void TypeDatabase::add_argument(Method method, int n, const TypeSet type) {
   arguments->set(n, type);
 }
 
-void TypeDatabase::add_usage(int position, const TypeSet type) {
-  ASSERT(usage_.find(position) == usage_.end());
+void TypeDatabase::add_input(int position, int n, int size, const TypeSet type) {
+  TypeStack* stack = null;
+  if (n == 0) {
+    ASSERT(input_.find(position) == input_.end());
+    stack = new TypeStack(size - 1, size, words_per_type_);
+    input_[position] = stack;
+  } else {
+    stack = input_.at(position);
+  }
+  stack->set(n, type);
+}
+
+void TypeDatabase::add_output(int position, const TypeSet type) {
+  ASSERT(output_.find(position) == output_.end());
   TypeSet copy = copy_type(type);
   uint8 opcode = *(program_->bcp_from_absolute_bci(position));
-  usage_.emplace(position, copy);
+  output_.emplace(position, copy);
   returns_.emplace(position + opcode_length[opcode], copy);
 }
 

src/compiler/propagation/type_database.h

@@ -40,8 +40,7 @@ class TypeDatabase {
 
   const std::vector<Method> methods() const;
   const std::vector<TypeSet> arguments(Method method) const;
-  const TypeSet usage(int position) const;
-  const TypeSet return_type(int position) const;
+  const TypeSet output(int position) const;
 
   std::string as_json() const;
 
@@ -64,7 +63,9 @@ class TypeDatabase {
   std::vector<TypeStack*> types_;
 
   std::unordered_map<int, TypeStack*> methods_;
-  std::unordered_map<int, TypeSet> usage_;
+
+  std::unordered_map<int, TypeStack*> input_;
+  std::unordered_map<int, TypeSet> output_;
   std::unordered_map<int, TypeSet> returns_;
 
   static std::unordered_map<Program*, TypeDatabase*> cache_;
@@ -73,7 +74,8 @@ class TypeDatabase {
 
   void add_method(Method method);
   void add_argument(Method method, int n, const TypeSet type);
-  void add_usage(int position, const TypeSet type);
+  void add_input(int position, int n, int size, const TypeSet type);
+  void add_output(int position, const TypeSet type);
 
   TypeSet copy_type(const TypeSet type);
   TypeStack* add_types_block();

src/compiler/propagation/type_propagator.cc

@@ -259,16 +259,23 @@ void TypePropagator::propagate(TypeDatabase* types) {
     uint8* entry = program()->entry_task().entry();
     type.clear(words_per_type());
     type.add_instance(program()->task_class_id());
-    types->add_usage(program()->absolute_bci_from_bcp(entry), type);
+    types->add_output(program()->absolute_bci_from_bcp(entry), type);
   }
 
-  sites_.for_each([&](uint8* site, Set<TypeVariable*>& results) {
+  output_.for_each([&](uint8* site, Set<TypeVariable*>& output) {
     type.clear(words_per_type());
-    for (auto it = results.begin(); it != results.end(); it++) {
+    for (auto it = output.begin(); it != output.end(); it++) {
       type.add_all_also_blocks((*it)->type(), words_per_type());
     }
     int position = program()->absolute_bci_from_bcp(site);
-    types->add_usage(position, type);
+    types->add_output(position, type);
+  });
+
+  input_.for_each([&](uint8* site, std::vector<TypeVariable*>& variables) {
+    int position = program()->absolute_bci_from_bcp(site);
+    for (unsigned i = 0; i < variables.size(); i++) {
+      types->add_input(position, i, variables.size(), variables[i]->type());
+    }
   });
 
   // Group the methods and blocks based on the bytecode position, so
@@ -400,6 +407,7 @@ void TypePropagator::call_method(MethodTemplate* caller,
 void TypePropagator::call_static(MethodTemplate* caller, TypeScope* scope, uint8* site, Method target) {
   TypeStack* stack = scope->top();
   int arity = target.arity();
+  if (site) add_input(site, stack, arity);
 
   std::vector<ConcreteType> arguments;
   stack->push_empty();
@@ -470,6 +478,7 @@ void TypePropagator::call_static(MethodTemplate* caller, TypeScope* scope, uint8
 void TypePropagator::call_virtual(MethodTemplate* caller, TypeScope* scope, uint8* site, int arity, int offset) {
   TypeStack* stack = scope->top();
   TypeSet receiver = stack->local(arity - 1);
+  if (site) add_input(site, stack, arity);
 
   std::vector<ConcreteType> arguments;
   stack->push_empty();
@@ -501,6 +510,7 @@ void TypePropagator::call_block(TypeScope* scope, uint8* site, int arity) {
   TypeStack* stack = scope->top();
   TypeSet receiver = stack->local(arity - 1);
   BlockTemplate* block = receiver.block();
+  if (site) add_input(site, stack, arity);
 
   // If we're passing too few arguments to the block, this will
   // throw so we do not need to update the block's argument types.
@@ -568,8 +578,9 @@ void TypePropagator::load_lambda(TypeScope* scope, Method method) {
 
 void TypePropagator::load_field(MethodTemplate* user, TypeStack* stack, uint8* site, int index) {
   TypeSet instance = stack->local(0);
-  stack->push_empty();
+  if (site) add_input(site, stack, 1);
 
+  stack->push_empty();
   TypeSet::Iterator it(instance, words_per_type_);
   while (it.has_next()) {
     unsigned id = it.next();
@@ -603,7 +614,7 @@ void TypePropagator::load_outer(TypeScope* scope, uint8* site, int index) {
   // this particular access site. We use this merged type exclusively
   // for the output of the type propagator, so we don't actually
   // use the merged type anywhere in the analysis.
-  TypeVariable* merged = this->outer(site);
+  TypeVariable* merged = this->output(site);
   merged->type().add_all_also_blocks(value, words_per_type());
 }
 
@@ -623,7 +634,7 @@ bool TypePropagator::handle_typecheck_result(TypeScope* scope, uint8* site, bool
       scope->throw_maybe();
     }
   }
-  outer(site)->type().add_all(stack->local(0), words_per_type());
+  output(site)->type().add_all(stack->local(0), words_per_type());
   if (as_check) stack->pop();
   return can_succeed;
 }
@@ -657,13 +668,13 @@ TypeVariable* TypePropagator::global_variable(int index) {
   }
 }
 
-TypeVariable* TypePropagator::outer(uint8* site) {
+TypeVariable* TypePropagator::output(uint8* site) {
   auto it = outers_.find(site);
   if (it == outers_.end()) {
-    TypeVariable* variable = new TypeVariable(words_per_type());
-    outers_[site] = variable;
-    add_site(site, variable);
-    return variable;
+    TypeVariable* output = new TypeVariable(words_per_type());
+    outers_[site] = output;
+    add_output(site, output);
+    return output;
   } else {
     return it->second;
   }
@@ -675,8 +686,29 @@ void TypePropagator::enqueue(MethodTemplate* method) {
   enqueued_.push_back(method);
 }
 
-void TypePropagator::add_site(uint8* site, TypeVariable* result) {
-  sites_[site].insert(result);
+void TypePropagator::add_input(uint8* site, TypeStack* input, int n) {
+  auto probe = input_.find(site);
+  if (probe != input_.end()) {
+    std::vector<TypeVariable*>& variables = probe->second;
+    ASSERT(variables.size() == n);
+    for (int i = 0; i < n; i++) {
+      TypeVariable* variable = variables[i];
+      variable->type().add_all_also_blocks(input->local(n - i - 1), words_per_type());
+    }
+    return;
+  }
+
+  std::vector<TypeVariable*> variables;
+  for (int i = 0; i < n; i++) {
+    TypeVariable* variable = new TypeVariable(words_per_type());
+    variable->type().add_all_also_blocks(input->local(n - i - 1), words_per_type());
+    variables.push_back(variable);
+  }
+  input_[site] = variables;
+}
+
+void TypePropagator::add_output(uint8* site, TypeVariable* output) {
+  output_[site].insert(output);
 }
 
 MethodTemplate* TypePropagator::find_method(Method target, std::vector<ConcreteType> arguments) {

src/compiler/propagation/type_propagator.h

@@ -66,10 +66,12 @@ class TypePropagator {
 
   TypeVariable* global_variable(int index);
   TypeVariable* field(unsigned type, int index);
-  TypeVariable* outer(uint8* site);
+  TypeVariable* output(uint8* site);
 
   void enqueue(MethodTemplate* method);
-  void add_site(uint8* site, TypeVariable* result);
+
+  void add_input(uint8* site, TypeStack* input, int n);
+  void add_output(uint8* site, TypeVariable* output);
 
 #define ENSURE_ENTRY_POINT(name, symbol, arity) \
   void ensure_##name();
@@ -85,13 +87,14 @@ class TypePropagator {
   ENTRY_POINTS(HAS_ENTRY_POINT)
 #undef HAS_ENTRY_POINT
 
-  Map<uint8*, Set<TypeVariable*>> sites_;
+  Map<uint8*, std::vector<TypeVariable*>> input_;
+  Map<uint8*, Set<TypeVariable*>> output_;
 
   std::unordered_map<uint32, MethodTemplate*> methods_;
   std::unordered_map<uint32, BlockTemplate*> blocks_;
 
   std::unordered_map<int, TypeVariable*> globals_;
-  std::unordered_map<uint8*, TypeVariable*> outers_;
+  std::unordered_map<uint8*, TypeVariable*> outers_;  // TODO(kasper): Rename this.
   std::unordered_map<unsigned, std::unordered_map<int, TypeVariable*>> fields_;
   std::vector<MethodTemplate*> enqueued_;
 

src/compiler/propagation/type_stack.h

@@ -44,6 +44,8 @@ class TypeStack {
   }
 
   int sp() const { return sp_; }
+
+  int size() const { return size_; }
   int available() const { return size_ - (sp_ + 1); }
 
   TypeSet get(int index) {

src/compiler/propagation/type_variable.cc

@@ -20,7 +20,7 @@ namespace toit {
 namespace compiler {
 
 TypeSet TypeVariable::use(TypePropagator* propagator, MethodTemplate* user, uint8* site) {
-  if (site) propagator->add_site(site, this);
+  if (site) propagator->add_output(site, this);
   if (user) users_.insert(user);
   return type();
 }

tests/type_propagation/gold/array_do_test.gold

@@ -1,18 +1,18 @@
 main tests/type_propagation/array_do_test.toit
   0[025] - load smi 1
-  1[053] - invoke static create_array_ <sdk>/core/collections.toit // {LargeArray_|SmallArray_}
-  4[053] - invoke static create_list_literal_from_array_ <sdk>/core/collections.toit // {List_}
+  1[053] - invoke static create_array_ <sdk>/core/collections.toit // [{SmallInteger_}] -> {LargeArray_|SmallArray_}
+  4[053] - invoke static create_list_literal_from_array_ <sdk>/core/collections.toit // [{LargeArray_|SmallArray_}] -> {List_}
   7[029] - load [block] in main tests/type_propagation/array_do_test.toit
  12[015] - load local 1
  13[038] - load block 1
- 15[058] - invoke virtual do // {Null_}
+ 15[058] - invoke virtual do // [{List_}, [block]] -> {Null_}
  19[089] - return null S3 0
 
 [block] in main tests/type_propagation/array_do_test.toit
  - argument 0: [block]
  - argument 1: {*}
   0[016] - load local 2
-  1[053] - invoke static id tests/type_propagation/array_do_test.toit // {*}
+  1[053] - invoke static id tests/type_propagation/array_do_test.toit // [{*}] -> {*}
   4[088] - return S1 2
 
 id tests/type_propagation/array_do_test.toit

tests/type_propagation/gold/array_do_test.gold-O2

@@ -1,18 +1,18 @@
 main tests/type_propagation/array_do_test.toit
   0[025] - load smi 1
-  1[053] - invoke static create_array_ <sdk>/core/collections.toit // {LargeArray_|SmallArray_}
-  4[053] - invoke static create_list_literal_from_array_ <sdk>/core/collections.toit // {List_}
+  1[053] - invoke static create_array_ <sdk>/core/collections.toit // [{SmallInteger_}] -> {LargeArray_|SmallArray_}
+  4[053] - invoke static create_list_literal_from_array_ <sdk>/core/collections.toit // [{LargeArray_|SmallArray_}] -> {List_}
   7[029] - load [block] in main tests/type_propagation/array_do_test.toit
  12[015] - load local 1
  13[038] - load block 1
- 15[058] - invoke virtual do // {Null_}
+ 15[058] - invoke virtual do // [{List_}, [block]] -> {Null_}
  19[089] - return null S3 0
 
 [block] in main tests/type_propagation/array_do_test.toit
  - argument 0: [block]
  - argument 1: {*}
   0[016] - load local 2
-  1[053] - invoke static id tests/type_propagation/array_do_test.toit // {*}
+  1[053] - invoke static id tests/type_propagation/array_do_test.toit // [{*}] -> {*}
   4[088] - return S1 2
 
 id tests/type_propagation/array_do_test.toit