from __future__ import annotations
import asyncio
from collections.abc import Callable, Coroutine, Generator, Mapping, Sequence
from logging import getLogger
from pathlib import Path
import threading
from types import UnionType
from typing import cast, overload
from jitx._translate.lookup import ComputedNet
from jitx._utils.unionfind import UnionFind
from jitx.copper import Copper
from jitx.net import Net, Port, PortAttachment, TopologyNet
from jitx.component import Component
from jitx.landpattern import PadMapping
from jitx.stackup import Stackup
from jitx.via import Via
from jitxcore._proto import idmap_pb2, messages_pb2
from .._translate.idmap import ReverseMapper, fqnameob
from .._websocket import Envelope, Message, PersistentWebSocketClient, runtime_uri
from ..design import Design
from ..inspect import Trace, extract, visit
from ..query import query
from ..transform import IDENTITY, Transform
logger = getLogger(__name__)
def _warn_if_issues(rmap: ReverseMapper) -> None:
if not rmap.unattached_copper_warning:
return
W = 38
def row(text: str = "") -> str:
return "|" + (" " + text).center(W) + "|"
lines = [
"",
"+" + "-" * W + "+",
row("DESIGN COMPLETED WITH WARNINGS"),
row(),
row("The design may not export correctly."),
row("Please review the warning(s) above."),
"+" + "-" * W + "+",
"",
]
logger.warning("\n".join(lines))
[docs]
class Runtime:
class __Active(threading.local):
runtime: list[Runtime | SyncRuntime | None]
def __init__(self):
super().__init__()
self.runtime = []
active = __Active()
@overload
def __init__(self): ...
@overload
def __init__(self, *, uri: str): ...
@overload
def __init__(self, *, project: str | Path): ...
def __init__(self, uri: str | None = None, project: str | Path | None = None):
import jitx.run
# if a Runtime object is constructed, it's fairly safe to assume that
# we're not trying to autodetect a design run as '__main__'.
jitx.run.autodetect_design = False
self.__uri = uri or (runtime_uri(project=project) if project else runtime_uri())
self.__client = PersistentWebSocketClient(uri=self.__uri)
self.__reentered = 0
@property
def _client(self):
return self.__client
[docs]
def decouple(self):
return self.__class__(uri=self.__uri)
def __enter(self, r: Runtime | SyncRuntime):
Runtime.active.runtime.append(r)
if not self.__reentered:
self.__client.keepalive()
self.__reentered += 1
return self.__reentered == 1
def __exit(self):
self.__reentered -= 1
Runtime.active.runtime.pop()
return self.__reentered == 0
async def __aenter__(self) -> Runtime:
if self.__enter(self):
await self.open()
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
if self.__exit():
await self.close()
[docs]
async def open(self):
await self.__client.establish()
[docs]
async def close(self):
await self.__client.disconnect()
def __enter__(self) -> SyncRuntime:
sr = SyncRuntime(self)
if self.__enter(sr):
sr.open()
return sr
def __exit__(self, exc_type, exc_val, exc_tb):
sr = Runtime.active.runtime[-1]
if self.__exit():
# should always be true.
if isinstance(sr, SyncRuntime):
sr.close()
[docs]
@staticmethod
def require():
if Runtime.active.runtime:
r = Runtime.active.runtime[-1]
if not isinstance(r, Runtime | SyncRuntime):
raise TypeError("Active runtime is not a Runtime")
return r
else:
return Runtime()
def _construct[T: Design](self, design: T | type[T]) -> T:
return _instantiate_design(design)
[docs]
async def submit[T: Design](self, design: T | type[T], *, name: str | None = None):
constructed = self._construct(design)
packaged, rmap = _package_design(constructed)
design_name = name or fqnameob(constructed)
rd = AsyncRuntimeDesign(constructed, design_name, rmap, self)
conversation = await self.__client.root().request(
Message(
"des",
"request",
messages_pb2.Request(
design_name=design_name,
design_manager=messages_pb2.DesignManagerRequest(
load=messages_pb2.LoadDesignTrigger(design=packaged)
),
),
),
)
async for envelope in conversation:
async with envelope.process(namespace="des") as des:
@des
def response(response: messages_pb2.Response):
if response.HasField("design_manager"):
pass
# trigger = response.design_manager.trigger
# if trigger.message:
# result_message = trigger.message
elif response.HasField("error"):
raise RuntimeError(f"Load failed: {response.error.message}")
@des
def ids(idmap: idmap_pb2.IDMap):
from jitx._instantiation import instantiation
# flattening might generate IDs that generate access to implicit
# elements, which then needs to trigger instantiation.
with instantiation.require(), instantiation.frame():
rd._rmap.populate_stable_id_chains(idmap)
@des
async def stdin(data: Envelope):
await conversation.send(data.message.response({"message": "N"}))
@des
def stdout(data: Mapping):
logger.info(data["message"])
@des
def unhandled():
raise RuntimeError(
f"Unhandled response message: {envelope.type}" # noqa: B023
)
return rd
[docs]
async def capture[T: RuntimeDesign](
self, design: T, *, name: str | None = None
) -> T:
from jitx._translate.reverse_flow.linker import Capture
if name is None:
name = fqnameob(design.root)
conversation = await self.__client.route(f"design/{name}").request(
Message(
"des",
"request",
messages_pb2.Request(
design_name=name,
physical_design=messages_pb2.PhysicalDesignRequest(
get_layout_output=messages_pb2.GetLayoutOutputQuery()
),
),
)
)
async for envelope in conversation:
with envelope.process(namespace="des") as des:
@des
def response(response: messages_pb2.Response):
layout = response.physical_design.layout_output.output
linker = Capture(layout, design._rmap, design=design.root)
linker.run()
_warn_if_issues(design._rmap)
pass
@des
def stdout(data: Mapping):
logger.info(data["message"])
@des
def failure(data: Mapping):
raise RuntimeError(
f"Failed to capture output ({data['kind']}): {data['message']}"
)
@des
def unhandled():
raise RuntimeError(
f"Unhandled response message: {envelope.type}" # noqa: B023
)
design._invalidate()
return design
[docs]
class SyncRuntime:
class _AsyncBridgeThread:
def __init__(self):
self._loop: asyncio.AbstractEventLoop | None = None
self._thread: threading.Thread | None = None
self._ready = threading.Event()
def start(self) -> None:
if self._thread is not None:
return
def _run():
loop = asyncio.new_event_loop()
self._loop = loop
asyncio.set_event_loop(loop)
# link calling and running thread's runtime stack
# it's a blocking call, so it's safe to do so.
self._ready.set()
try:
loop.run_forever()
finally:
pending = asyncio.all_tasks(loop)
for t in pending:
t.cancel()
loop.run_until_complete(
asyncio.gather(*pending, return_exceptions=True)
)
loop.close()
self._thread = threading.Thread(
target=_run, name="jitx-sync-async-bridge", daemon=True
)
self._thread.start()
self._ready.wait()
def run(self, coro):
assert self._loop is not None, "loop thread not started"
return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
def close(self) -> None:
if self._loop is None or self._thread is None:
return
self._loop.call_soon_threadsafe(self._loop.stop)
self._thread.join()
self._loop = None
self._thread = None
self._ready.clear()
def __init__(self, runtime: Runtime):
self.__runtime = runtime
self.__loop_thread = SyncRuntime._AsyncBridgeThread()
[docs]
def runasync[T](self, coro: Coroutine[None, None, T]) -> T:
from jitx._instantiation import instantiation
self.__loop_thread.start()
outer_instantiation = instantiation.data.stack
async def decoupler():
instantiation.data.stack = outer_instantiation
async with self.__runtime.decouple():
return await coro
return self.__loop_thread.run(decoupler())
# Technically we'd be free to run an event loop here if there isn't one,
# however, if we, for example, run websockets open like this, and then
# websockets close the same way, those will be in different event loops,
# and it has callbacks from the first event loop running in the second
# which it does not like.
# return asyncio.run(coro)
# So it turns out it doesn't matter whether there's an event loop on
# this thread or not, we always have to start a new one.
def __enter__(self):
return self.__runtime.__enter__()
def __exit__(self, exc_type, exc_ob, exc_tb):
return self.__runtime.__exit__(exc_type, exc_ob, exc_tb)
# these should logically not be possible to construct from code,
# and thus could just be exceptions, but it seems fine to keep the parity.
# having them defined allows the Runtime|SyncRuntime return type to not
def __aenter__(self):
return self.__runtime.__aenter__()
def __aexit__(self, exc_type, exc_ob, exc_tb):
return self.__runtime.__aexit__(exc_type, exc_ob, exc_tb)
[docs]
def close(self):
self.runasync(self.__runtime.close())
self.__loop_thread.close()
[docs]
def open(self):
self.runasync(self.__runtime.open())
[docs]
def submit[T: Design](self, design: T | type[T], *, name: str | None = None):
ard = self.runasync(self.__runtime.submit(design, name=name))
return SyncRuntimeDesign(ard.root, ard.name, ard._rmap, self)
[docs]
def capture[T: RuntimeDesign](self, design: T, *, name: str | None = None) -> T:
return self.runasync(self.__runtime.capture(design, name=name))
class _RuntimeDispatch:
@overload
def __call__(self): ...
@overload
def __call__(self, *, uri: str): ...
@overload
def __call__(self, *, project: str | Path): ...
def __call__(self, *, uri: str | None = None, project: str | Path | None = None):
if uri:
r = Runtime(uri=uri)
elif project:
r = Runtime(project=project)
else:
r = Runtime.require()
return r
def __enter__(self):
return self().__enter__()
def __exit__(self, exc_type, exc_ob, exc_tb):
r = Runtime.active.runtime[-1]
assert r is not None
return r.__exit__(exc_type, exc_ob, exc_tb)
async def __aenter__(self):
return await self().__aenter__()
async def __aexit__(self, exc_type, exc_ob, exc_tb):
r = Runtime.active.runtime[-1]
assert r is not None
return await r.__aexit__(exc_type, exc_ob, exc_tb)
async def submit[T: Design](self, design: T | type[T], *, name: str | None = None):
async with Runtime.require() as r:
return await r.submit(design, name=name)
async def capture[T: RuntimeDesign](
self, design: T, *, name: str | None = None
) -> T:
async with Runtime.require() as r:
return await r.capture(design, name=name)
@staticmethod
def require():
return Runtime.require()
dispatch = _RuntimeDispatch()
[docs]
class RuntimeDesign[T: Design]:
root: T
_rmap: ReverseMapper
def __init__(self, design: T, name: str, rmap: ReverseMapper):
self.root = design
self._rmap = rmap
self.__name = name
self.__memonets: RuntimeDesign.Nets | None = None
def _invalidate(self):
self.__memonets = None
@property
def name(self):
return self.__name
[docs]
def query[Q](
self,
target: type[Q] | tuple[type[Q], ...],
/,
*,
through: tuple[type, ...] | UnionType | None = None,
transform: Transform | None = IDENTITY,
opaque: tuple[type, ...] | UnionType | None = None,
refs: bool = False,
filter: Callable[[Q], bool] | None = None,
) -> Generator[tuple[Trace, Q], None, None]:
"""Convenience method to query this design. See
:py:meth:`~jitx.query.TransformQuery.query` for details."""
return query(
self.root,
target,
through=through,
transform=transform,
opaque=opaque,
refs=refs,
filter=filter,
)
[docs]
def nets(self):
if self.__memonets is not None:
return self.__memonets
uf: UnionFind[Net | TopologyNet] = UnionFind()
for net in extract(self.root, Net | TopologyNet):
uf.add(net)
if cn := ComputedNet.get(net):
if cn.net:
uf.union(net, cn.net)
non_computed_count = 0
for net in extract(self.root, Net | TopologyNet):
if isinstance(net, Net):
def recurse(net):
count = 0
for n in net._connected:
if isinstance(n, Net):
if not uf.connected(net, n):
count += 1
uf.union(net, n)
count += recurse(n)
return count
non_computed_count += recurse(net)
if non_computed_count:
logger.debug(
"Unioned nets without computed net connection: %d", non_computed_count
)
missing_nets: list[tuple[Trace, Component, set[Port]]] = []
for trace, comp in visit(self.root, (Component, PortAttachment)):
if isinstance(comp, Component):
missing_net_on_comp: list[Port] = []
for pm in extract(comp, PadMapping):
for port, pads in pm.items():
cn = ComputedNet.get(port)
if not cn:
missing_net_on_comp.append(port)
continue
if not isinstance(pads, Sequence):
pads = (pads,)
for pad in pads:
cn.assign(pad)
if missing_net_on_comp:
missing_nets.append((trace, comp, set(missing_net_on_comp)))
elif isinstance(comp, PortAttachment):
if isinstance(comp.port, Sequence):
for ap, cp in zip(
extract(comp.attachment, Port, filter=Port.is_single_pin),
comp.port,
strict=False,
):
cn = ComputedNet.get(cp)
if cn:
cn.assign(ap)
else:
cn = ComputedNet.get(comp.port)
if cn:
cn.assign(comp.attachment)
if missing_nets:
for trace, comp, ports in missing_nets:
portnames = []
for t, _pt in visit(
comp, Port, filter=lambda p, ports=ports: p in ports
):
portnames.append(str(t.path))
if len(portnames) > 1:
portnames.insert(0, "") # don't leave one on the same line
logger.warning(
"no computed net for port in component %s: %s",
trace.path,
"\n ".join(portnames),
)
self.__memonets = self.Nets(uf)
return self.__memonets
[docs]
def layers(self):
return self.Layers(self.root.substrate.stackup)
[docs]
class Net[N: Port]:
def __init__(self, nets: Sequence[Net[N] | TopologyNet[N]]):
cname = None
for net in nets:
nname = net._connected_name
if nname:
if cname and cname != nname:
logger.warning(
"Duplicate connected names encountered for nets '%s' and '%s'",
cname,
net._connected_name,
)
else:
cname = nname
self.nets = nets
self.name = cname
@overload
def __iter__(
self: RuntimeDesign.Net[Port],
) -> Generator[Port | Copper | Via, None, None]: ...
@overload
def __iter__(self) -> Generator[N, None, None]: ...
def __iter__(self):
# cast to erase the port type, or pyright thinks the overloads
# might be incompatible - they can't really be at this point,
# but hard to prove.
nets = cast(Sequence[Net], self.nets)
for net in nets:
yield from net
[docs]
class Nets:
def __init__(self, uf: UnionFind[Net | TopologyNet]):
self.__uf = uf
self.__nets = {
root: RuntimeDesign.Net(group) for root, group in uf.groups().items()
}
[docs]
def find(self, element) -> RuntimeDesign.Net | None:
if isinstance(element, Net | TopologyNet):
net = element
else:
cnet = ComputedNet.get(element)
if cnet and cnet.net:
net = cnet.net
else:
return None
if net not in self.__uf:
return None
return self.__nets.get(self.__uf.find(net))
[docs]
class Layers:
def __init__(self, stackup: Stackup):
self.__conductor_count = len(stackup.conductors)
[docs]
def normalize(self, layer: int):
if layer < 0:
layer += self.__conductor_count
return layer
[docs]
class AsyncRuntimeDesign[T: Design](RuntimeDesign[T]):
def __init__(self, design: T, name: str, rmap: ReverseMapper, runtime: Runtime):
super().__init__(design, name, rmap)
self.__runtime = runtime
[docs]
async def capture(self):
return await self.__runtime.capture(self)
[docs]
class SyncRuntimeDesign[T: Design](RuntimeDesign[T]):
def __init__(self, design: T, name: str, rmap: ReverseMapper, runtime: SyncRuntime):
super().__init__(design, name, rmap)
self.__runtime = runtime
[docs]
def capture(self):
return self.__runtime.capture(self)
def _instantiate_design[T: Design](design: T | type[T]) -> T:
from jitx._instantiation import instantiation
from jitx._structural import Instantiable, Proxy
import gc
if isinstance(design, Instantiable):
cls = design._instantiable_()
if not (isinstance(cls, type) and issubclass(cls, Design)):
raise ValueError(f"Not an instantiable design: {cls}")
with instantiation.require():
with instantiation.frame():
instantiated = design._instantiate_({})
elif isinstance(design, type):
if not issubclass(design, Design):
raise ValueError(f"Not a Design subclass: {cls}")
try:
if hasattr(design, "__signature__"):
design.__signature__.bind()
except TypeError:
raise ValueError(
"Design is parameterized but no parameters were provided"
) from None
with instantiation.require():
with instantiation.frame():
instantiated = design()
else:
instantiated = design
assert isinstance(instantiated, Design), (
f"Instantiation resulted in non-Design object {Proxy.type(instantiated)}"
)
gc.collect()
return cast(T, instantiated)
def _package_design(design: Design):
from jitx._instantiation import instantiation
from jitx._translate.design import package_design
from .._translate.idmap import idmap
if not isinstance(design, Design):
raise TypeError(f"{design} is not a Design class")
idmap.clear()
with instantiation.require():
packaged = package_design(design)
rmap = idmap.export()
return packaged, rmap
