Refmap optimization

edg/core/Array.py

@@ -24,8 +24,9 @@ def __init__(self, container: Vector, elt: ConstraintExpr):
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return [self.container]
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
-    contained_map = self.container._get_contained_ref_map()
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
+    contained_map = self.container._elt_sample._create_ref_map(edgir.LocalPath())
+
     pb = edgir.ValueExpr()
     pb.map_extract.container.ref.CopyFrom(ref_map[self.container])  # TODO support arbitrary refs
     pb.map_extract.path.CopyFrom(contained_map[self.elt])
@@ -40,7 +41,7 @@ def __init__(self, elts: ConstraintExpr):
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return [self.elts]
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.unary_set.op = edgir.UnarySetExpr.Op.FLATTEN
     pb.unary_set.vals.CopyFrom(self.elts._expr_to_proto(ref_map))
@@ -178,16 +179,15 @@ def _instance_to_proto(self) -> edgir.PortLike:
   def _def_to_proto(self) -> edgir.PortTypes:
     raise RuntimeError()  # this doesn't generate into a library element
 
-  def _get_ref_map(self, prefix: edgir.LocalPath) -> IdentityDict[Refable, edgir.LocalPath]:
+  def _build_ref_map(self, ref_map: Refable.RefMapType, prefix: edgir.LocalPath) -> None:
+    super()._build_ref_map(ref_map, prefix)
+    ref_map[self._length] = edgir.localpath_concat(prefix, edgir.LENGTH)
+    ref_map[self._requested] = edgir.localpath_concat(prefix, edgir.ALLOCATED)
     elts_items = self._elts.items() if self._elts is not None else []
-
-    return super()._get_ref_map(prefix) + IdentityDict[Refable, edgir.LocalPath](
-      [(self._length, edgir.localpath_concat(prefix, edgir.LENGTH)),
-       (self._requested, edgir.localpath_concat(prefix, edgir.ALLOCATED))],
-      *[elt._get_ref_map(edgir.localpath_concat(prefix, index)) for (index, elt) in elts_items]) + IdentityDict(
-      *[elt._get_ref_map(edgir.localpath_concat(prefix, edgir.Allocate(suggested_name)))
-        for (suggested_name, elt) in self._requests]
-    )
+    for index, elt in elts_items:
+      elt._build_ref_map(ref_map, edgir.localpath_concat(prefix, index))
+    for suggested_name, request in self._requests:
+      request._build_ref_map(ref_map, edgir.localpath_concat(prefix, edgir.Allocate(suggested_name)))
 
   def _get_initializers(self, path_prefix: List[str]) -> List[Tuple[ConstraintExpr, List[str], ConstraintExpr]]:
     if self._elts is not None:
@@ -196,9 +196,6 @@ def _get_initializers(self, path_prefix: List[str]) -> List[Tuple[ConstraintExpr
     else:
       return []
 
-  def _get_contained_ref_map(self) -> IdentityDict[Refable, edgir.LocalPath]:
-    return self._elt_sample._get_ref_map(edgir.LocalPath())
-
   def defined(self) -> None:
     """Marks this vector as defined, even if it is empty. Can be called multiple times, and append_elt can continue
     to be used.

edg/core/ArrayExpr.py

@@ -20,7 +20,7 @@ def __init__(self):
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:  # element should be returned by the containing ConstraintExpr
     return []
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     raise ValueError  # can't be used directly
 
 

edg/core/Binding.py

@@ -73,7 +73,7 @@ def is_bound(self) -> bool:
     return True
 
   @abstractmethod
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr: ...
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr: ...
 
   @abstractmethod
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]: ...
@@ -100,7 +100,7 @@ def is_bound(self) -> bool:
     else:
       return True
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.ref.CopyFrom(ref_map[expr])
     return pb
@@ -128,7 +128,7 @@ def __init__(self, value: bool):
     super().__init__()
     self.value = value
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.literal.boolean.val = self.value
     return pb
@@ -141,7 +141,7 @@ def __repr__(self) -> str:
   def __init__(self, value: int):
     self.value = value
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.literal.integer.val = self.value
     return pb
@@ -154,7 +154,7 @@ def __repr__(self) -> str:
   def __init__(self, value: Union[float, int]):
     self.value: float = float(value)
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.literal.floating.val = self.value
     return pb
@@ -167,7 +167,7 @@ def __repr__(self) -> str:
   def __init__(self, value: Range):
     self.value = value
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.literal.range.minimum.floating.val = self.value.lower
     pb.literal.range.maximum.floating.val = self.value.upper
@@ -182,7 +182,7 @@ def __init__(self, value: str):
     super().__init__()
     self.value = value
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.literal.text.val = self.value
     return pb
@@ -196,7 +196,7 @@ def __init__(self, values: Sequence[LiteralBinding]):
     super().__init__()
     self.values = values
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.literal.array.SetInParent()
     for value in self.values:
@@ -218,7 +218,7 @@ def __init__(self, lower: FloatExpr, upper: FloatExpr):
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return chain(self.lower._get_exprs(), self.lower._get_exprs())
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.binary.op = edgir.BinaryExpr.RANGE
     pb.binary.lhs.CopyFrom(self.lower._expr_to_proto(ref_map))
@@ -234,7 +234,7 @@ def __init__(self, values: Sequence[ConstraintExpr]):
     super().__init__()
     self.values = values
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.array.SetInParent()
     for value in self.values:
@@ -267,7 +267,7 @@ def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return chain(self.src._get_exprs())
 
   def expr_to_proto(self, expr: ConstraintExpr,
-                    ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+                    ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.unary.op = self.op_map[self.op]
     pb.unary.val.CopyFrom(self.src._expr_to_proto(ref_map))
@@ -303,7 +303,7 @@ def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return chain(self.src._get_exprs())
 
   def expr_to_proto(self, expr: ConstraintExpr,
-                    ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+                    ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.unary_set.op = self.op_map[self.op]
     pb.unary_set.vals.CopyFrom(self.src._expr_to_proto(ref_map))
@@ -352,7 +352,7 @@ def __init__(self,
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return chain(self.lhs._get_exprs(), self.rhs._get_exprs())
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.binary.op = self.op_map[self.op]
     pb.binary.lhs.CopyFrom(self.lhs._expr_to_proto(ref_map))
@@ -381,7 +381,7 @@ def __init__(self,
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return chain(self.lhset._get_exprs(), self.rhs._get_exprs())
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.binary_set.op = self.op_map[self.op]
     pb.binary_set.lhset.CopyFrom(self.lhset._expr_to_proto(ref_map))
@@ -402,7 +402,7 @@ def __init__(self, cond: BoolExpr, then_val: ConstraintExpr, else_val: Constrain
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return chain(self.cond._get_exprs(), self.then_val._get_exprs(), self.else_val._get_exprs())
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.ifThenElse.cond.CopyFrom(self.cond._expr_to_proto(ref_map))
     pb.ifThenElse.tru.CopyFrom(self.then_val._expr_to_proto(ref_map))
@@ -421,7 +421,7 @@ def __init__(self, src: Port):
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return [self.src]
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.ref.CopyFrom(ref_map[self.src])
     pb.ref.steps.add().reserved_param = edgir.IS_CONNECTED
@@ -438,7 +438,7 @@ def __init__(self, src: Union[BaseBlock, BasePort]):
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return []
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.ref.CopyFrom(ref_map[self.src])
     pb.ref.steps.add().reserved_param = edgir.NAME
@@ -456,7 +456,7 @@ def __init__(self, src: Vector):
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return [self.src]
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.ref.CopyFrom(ref_map[self.src])
     pb.ref.steps.add().reserved_param = edgir.LENGTH
@@ -474,7 +474,7 @@ def __init__(self, src: Vector):
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return [self.src]
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.ref.CopyFrom(ref_map[self.src])
     pb.ref.steps.add().reserved_param = edgir.ALLOCATED
@@ -484,7 +484,7 @@ def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edg
 class AssignBinding(Binding):
   # Convenience method to make an assign expr without the rest of this proto infrastructure
   @staticmethod
-  def make_assign(target: ConstraintExpr, value: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def make_assign(target: ConstraintExpr, value: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     pb = edgir.ValueExpr()
     pb.assign.dst.CopyFrom(ref_map[target])
     pb.assign.src.CopyFrom(value._expr_to_proto(ref_map))
@@ -501,5 +501,5 @@ def __init__(self, target: ConstraintExpr, value: ConstraintExpr):
   def get_subexprs(self) -> Iterable[Union[ConstraintExpr, BasePort]]:
     return [self.value]
 
-  def expr_to_proto(self, expr: ConstraintExpr, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def expr_to_proto(self, expr: ConstraintExpr, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     return self.make_assign(self.target, self.value, ref_map)

edg/core/Blocks.py

@@ -180,9 +180,14 @@ def make_connection(self) -> Optional[Union[ConnectedLink, Export]]:
     bridged_connects: List[Tuple[BasePort, edgir.LocalPath]] = []
     link_connects: List[Tuple[BasePort, edgir.LocalPath]] = []
     assert self.link_instance is not None
-    link_ref_map = self.link_instance._get_ref_map_allocate(edgir.LocalPath())
+    link_ref_map = self.link_instance._create_ref_map()
     for link_port, connected_ports in self.link_connected_ports.items():
       link_ref = link_ref_map[link_port]
+      if isinstance(link_port, BaseVector):
+        allocate_path = edgir.LocalPath()
+        allocate_path.CopyFrom(link_ref)
+        allocate_path.steps.append(edgir.LocalStep(allocate=''))
+        link_ref = allocate_path
       for connected_port in connected_ports:
         bridged_port = self.bridged_ports.get(connected_port, None)
         if bridged_port is None:  # direct connection, no bridge
@@ -214,14 +219,13 @@ class DescriptionString():
   def __init__(self, *elts: Union[str, DescriptionStringElts]):
     self.elts = elts
 
-  def add_to_proto(self, pb, ref_map):
+  def add_to_proto(self, pb: edgir.BlockLikeTypes, ref_map: Refable.RefMapType):
     for elt in self.elts:
       if isinstance(elt, DescriptionStringElts):
         elt.set_elt_proto(pb, ref_map)
       elif isinstance(elt, str):
         new_phrase = pb.description.add()
         new_phrase.text = elt
-    return pb
 
   class FormatUnits(DescriptionStringElts):
     ref: ConstraintExpr
@@ -310,7 +314,7 @@ def _elaborated_def_to_proto(self) -> BaseBlockEdgirType:
 
     return self._def_to_proto()
 
-  def _populate_def_proto_block_base(self, pb: BaseBlockEdgirType) -> BaseBlockEdgirType:
+  def _populate_def_proto_block_base(self, pb: BaseBlockEdgirType) -> None:
     """Populates the structural parts of a block proto: parameters, ports, superclasses"""
     assert self._elaboration_state == BlockElaborationState.post_contents or \
            self._elaboration_state == BlockElaborationState.post_generate
@@ -343,7 +347,7 @@ def _populate_def_proto_block_base(self, pb: BaseBlockEdgirType) -> BaseBlockEdg
     for (name, port) in self._ports.items():
       edgir.add_pair(pb.ports, name).CopyFrom(port._instance_to_proto())
 
-    ref_map = self._get_ref_map(edgir.LocalPath())  # TODO dedup ref_map
+    ref_map = self._create_ref_map()
     for (name, port) in self._ports.items():
       if port in self._required_ports:
         if isinstance(port, Port):
@@ -359,58 +363,45 @@ def _populate_def_proto_block_base(self, pb: BaseBlockEdgirType) -> BaseBlockEdg
 
     self._constraints.finalize()  # needed for source locator generation
 
-    ref_map = self._get_ref_map(edgir.LocalPath())
     self._populate_metadata(pb.meta, self._metadata, ref_map)
 
-    return pb
-
-  def _populate_def_proto_port_init(self, pb: BaseBlockEdgirType) -> BaseBlockEdgirType:
-    ref_map = self._get_ref_map(edgir.LocalPath())  # TODO dedup ref_map
-
+  def _populate_def_proto_port_init(self, pb: BaseBlockEdgirType, ref_map: Refable.RefMapType) -> None:
     for (name, port) in self._ports.items():
       for (param, path, initializer) in port._get_initializers([name]):
         edgir.add_pair(pb.constraints, f"(init){'.'.join(path)}").CopyFrom(
           AssignBinding.make_assign(param, param._to_expr_type(initializer), ref_map)
         )
-    return pb
 
-  def _populate_def_proto_param_init(self, pb: BaseBlockEdgirType) -> BaseBlockEdgirType:
-    ref_map = self._get_ref_map(edgir.LocalPath())  # TODO dedup ref_map
+  def _populate_def_proto_param_init(self, pb: BaseBlockEdgirType, ref_map: Refable.RefMapType) -> None:
     for (name, param) in self._parameters.items():
       if param.initializer is not None:
         edgir.add_pair(pb.constraints, f'(init){name}').CopyFrom(
           AssignBinding.make_assign(param, param.initializer, ref_map)
         )
-    return pb
 
-  def _populate_def_proto_block_contents(self, pb: BaseBlockEdgirType) -> BaseBlockEdgirType:
+  def _populate_def_proto_block_contents(self, pb: BaseBlockEdgirType, ref_map: Refable.RefMapType) -> None:
     """Populates the contents of a block proto: constraints"""
     assert self._elaboration_state == BlockElaborationState.post_contents or \
            self._elaboration_state == BlockElaborationState.post_generate
 
     self._constraints.finalize()
 
-    ref_map = self._get_ref_map(edgir.LocalPath())
-
     for (name, constraint) in self._constraints.items():
       edgir.add_pair(pb.constraints, name).CopyFrom(constraint._expr_to_proto(ref_map))
 
-    return pb
-
-  def _populate_def_proto_description(self, pb: BaseBlockEdgirType) -> BaseBlockEdgirType:
+  def _populate_def_proto_description(self, pb: BaseBlockEdgirType, ref_map: Refable.RefMapType) -> None:
     description = self.description
     assert(description is None or isinstance(description, DescriptionString))
     if isinstance(description, DescriptionString):
-      pb = description.add_to_proto(pb, self._get_ref_map(edgir.LocalPath()))
-
-    return pb
-
-  def _get_ref_map(self, prefix: edgir.LocalPath) -> IdentityDict[Refable, edgir.LocalPath]:
-    return super()._get_ref_map(prefix) + IdentityDict(
-      [(self.name(), edgir.localpath_concat(prefix, edgir.NAME))],
-      *[param._get_ref_map(edgir.localpath_concat(prefix, name)) for (name, param) in self._parameters.items()],
-      *[port._get_ref_map(edgir.localpath_concat(prefix, name)) for (name, port) in self._ports.items()]
-    )
+      description.add_to_proto(pb, ref_map)
+
+  def _build_ref_map(self, ref_map: Refable.RefMapType, prefix: edgir.LocalPath) -> None:
+    super()._build_ref_map(ref_map, prefix)
+    ref_map[self.name()] = edgir.localpath_concat(prefix, edgir.NAME)
+    for name, param in self._parameters.items():
+      param._build_ref_map(ref_map, edgir.localpath_concat(prefix, name))
+    for name, port in self._ports.items():
+      port._build_ref_map(ref_map, edgir.localpath_concat(prefix, name))
 
   def _bind_in_place(self, parent: Union[BaseBlock, Port]):
     self._parent = parent

edg/core/ConstraintExpr.py

@@ -91,7 +91,7 @@ def _bind(self: SelfType, binding: Binding) -> SelfType:
   def _is_bound(self) -> bool:
     return self.binding is not None and self.binding.is_bound()
 
-  def _expr_to_proto(self, ref_map: IdentityDict[Refable, edgir.LocalPath]) -> edgir.ValueExpr:
+  def _expr_to_proto(self, ref_map: Refable.RefMapType) -> edgir.ValueExpr:
     assert self.binding is not None
     return self.binding.expr_to_proto(self, ref_map)