schema_solve.py 39 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678
  1. from .utilities import (prRed, prGreen, prPurple, prWhite,
  2. prOrange,
  3. wrapRed, wrapGreen, wrapPurple, wrapWhite,
  4. wrapOrange,)
  5. from .utilities import init_connections, init_dependencies, get_link_in_out
  6. from .base_definitions import SchemaUINode, custom_props_types, MantisNodeGroup, replace_types
  7. from .node_container_common import setup_custom_props_from_np
  8. # a class that solves Schema nodes
  9. from bpy.types import NodeGroupInput, NodeGroupOutput
  10. class SchemaSolver:
  11. def __init__(self, schema_dummy, nodes, prototype, signature=None,):
  12. self.all_nodes = nodes # this is the parsed tree from Mantis
  13. self.node = schema_dummy
  14. self.tree = prototype.node_tree
  15. self.uuid = self.node.uuid
  16. if signature:
  17. self.signature = signature
  18. else:
  19. self.signature = self.node.signature
  20. self.schema_nodes={}
  21. self.solved_nodes = {}
  22. self.incoming_connections = {}
  23. self.outgoing_connections = {}
  24. self.constant_in = {}
  25. self.constant_out = {}
  26. self.array_input_connections = {}
  27. self.array_output_connections = {}
  28. self.nested_schemas = {}
  29. self.autogenerated_nodes = {}
  30. self.held_links = []
  31. self.tree_path_names = [*self.node.signature] # same tree as the schema node
  32. self.autogen_path_names = ['SCHEMA_AUTOGENERATED', *self.node.signature[1:]]
  33. self.is_node_group = False
  34. if self.node.prototype.bl_idname == "MantisNodeGroup":
  35. self.is_node_group = True
  36. if self.node.inputs['Schema Length'].links:
  37. self.index_link = self.node.inputs['Schema Length'].links[0]
  38. else:
  39. self.index_link = None
  40. # this should never happen, but it's OK to check.
  41. if self.node.inputs["Schema Length"].is_linked:
  42. if (other := self.node.inputs['Schema Length'].links[0].from_node).prepared == False:
  43. raise RuntimeError(f"Schema Length cannot be determined for {self.node} because {other} is not prepared.")
  44. self.solve_length = self.node.evaluate_input("Schema Length")
  45. # I'm making this a property of the solver because the solver's data is modified as it solves each iteration
  46. self.index = 0
  47. self.init_schema_links()
  48. self.set_index_strings()
  49. # Sort the multi-input nodes in reverse order of ID, this ensures that they are
  50. # read in the order they were created
  51. for inp in self.node.inputs.values():
  52. inp.links.sort(key=lambda a : -a.multi_input_sort_id)
  53. for ui_node in self.tree.nodes:
  54. # first we need to fill the parameters of the schema nodes.
  55. # we use the bl_idname because all schema nodes should be single-instance
  56. signature = (*self.tree_path_names, ui_node.bl_idname)
  57. if isinstance(ui_node, SchemaUINode):
  58. # We use the schema node's "natural signature" here because it represents
  59. # the "original" signature of the schema UI group node since this schema
  60. # solver may be in a nested schema, and its node's signature may have
  61. # uuid/index attached.
  62. get_sig = (*self.node.ui_signature, ui_node.bl_idname)
  63. if not (mantis_node := self.all_nodes.get(get_sig)): raise RuntimeError(wrapRed(f"Not found: {get_sig}"))
  64. self.schema_nodes[signature] = mantis_node
  65. mantis_node.fill_parameters(ui_node)
  66. # HACK to make Group Nodes work
  67. if ui_node.bl_idname == "NodeGroupInput":
  68. from .schema_containers import SchemaConstInput
  69. mantis_node = SchemaConstInput(signature=signature, base_tree=self.node.base_tree, parent_schema_node=self.node)
  70. self.schema_nodes[signature] = mantis_node
  71. mantis_node.fill_parameters(ui_node)
  72. if ui_node.bl_idname == "NodeGroupOutput":
  73. from .schema_containers import SchemaConstOutput
  74. mantis_node = SchemaConstOutput(signature=signature, base_tree=self.node.base_tree, parent_schema_node=self.node)
  75. self.schema_nodes[signature] = mantis_node
  76. mantis_node.fill_parameters(ui_node)
  77. def set_index_strings(self):
  78. self.index_str = lambda : '.'+str(self.uuid)+'.'+str(self.index).zfill(4)
  79. self.prev_index_str = lambda : '.'+str(self.uuid)+'.'+str(self.index-1).zfill(4)
  80. if self.is_node_group:
  81. self.index_str=lambda : ''
  82. self.prev_index_str=lambda : ''
  83. def init_schema_links(self,):
  84. """ Sort and store the links to/from the Schema group node."""
  85. for item in self.tree.interface.items_tree:
  86. if item.item_type == 'PANEL': continue
  87. parent_name='Constant'
  88. if item.parent.name != '': # in an "prphan" item this is left blank , it is not None or an AttributeError.
  89. parent_name = item.parent.name
  90. match parent_name:
  91. case 'Connection':
  92. if item.in_out == 'INPUT':
  93. if incoming_links := self.node.inputs[item.identifier].links:
  94. self.incoming_connections[item.name] = incoming_links[0]
  95. else:
  96. self.incoming_connections[item.name] = None
  97. else: # OUTPUT
  98. if outgoing_links := self.node.outputs[item.identifier].links:
  99. self.outgoing_connections[item.name] = outgoing_links.copy()
  100. else:
  101. self.outgoing_connections[item.name] = []
  102. case 'Constant':
  103. if item.in_out == 'INPUT':
  104. if constant_in_links := self.node.inputs[item.identifier].links:
  105. self.constant_in[item.name] = constant_in_links[0]
  106. else:
  107. self.constant_in[item.name] = None
  108. else: # OUTPUT
  109. if constant_out_links := self.node.outputs[item.identifier].links:
  110. self.constant_out[item.name] = constant_out_links.copy()
  111. else:
  112. self.constant_out[item.name] = []
  113. case 'Array':
  114. if item.in_out == 'INPUT':
  115. if item.identifier not in self.array_input_connections.keys():
  116. self.array_input_connections[item.identifier]=[]
  117. if in_links := self.node.inputs[item.identifier].links:
  118. self.array_input_connections[item.identifier]=in_links.copy()
  119. # I am tempted to put a check here, but it is sufficient to
  120. # rely on hierarchy links ensuring the arrays are prepared.
  121. # and testing here for un-prepared arrays will mess up schema
  122. # relationships in non-hierarchy situations.
  123. # Still, I'd like to have an easier time catching these problems.
  124. else: # OUTPUT
  125. if item.identifier not in self.array_output_connections.keys():
  126. self.array_output_connections[item.identifier]=[]
  127. if out_links := self.node.outputs[item.identifier].links:
  128. self.array_output_connections[item.identifier] = out_links.copy()
  129. def gen_solve_iteration_mantis_nodes(self, frame_mantis_nodes, unprepared):
  130. for prototype_ui_node in self.tree.nodes:
  131. mantis_node_name = prototype_ui_node.name
  132. index_str = self.index_str()
  133. mContext=self.node.mContext
  134. if isinstance(prototype_ui_node, SchemaUINode):
  135. continue # IGNORE the schema interface nodes, we already made them in __init__()
  136. # they are reused for each iteration.
  137. elif prototype_ui_node.bl_idname in ['NodeFrame', 'NodeReroute']:
  138. continue # IGNORE stuff that is purely UI - frames, reroutes.
  139. elif prototype_ui_node.bl_idname in ['NodeGroupInput', 'NodeGroupOutput']:
  140. continue # we converted these to Schema Nodes because they represent a Group input.
  141. signature = (*self.autogen_path_names, mantis_node_name+index_str)
  142. ui_signature=(*self.signature, mantis_node_name)
  143. prototype_mantis_node = self.all_nodes[ui_signature]
  144. # the prototype_mantis_node was generated inside the schema when we parsed the tree.
  145. # it is the prototype of the mantis node which we make for this iteration
  146. # for Schema sub-nodes ... they need a prototype to init.
  147. if prototype_mantis_node.node_type in ['DUMMY', 'DUMMY_SCHEMA']:
  148. # We stored the prototype ui_node when creating the Mantis node.
  149. ui_node = prototype_mantis_node.prototype
  150. # if prototype_mantis_node is a group or schema: TODO changes are needed elsewhere to make this easier to read. LEGIBILITY
  151. if ui_node.bl_idname in ["MantisNodeGroup", "MantisSchemaGroup"]:
  152. mantis_node = prototype_mantis_node.__class__(
  153. signature, prototype_mantis_node.base_tree, prototype=ui_node,
  154. ui_signature = prototype_mantis_node.signature)
  155. # now let's copy the links from the prototype node
  156. if ui_node.bl_idname in ["MantisNodeGroup"]:
  157. mantis_node.prepared = False
  158. mantis_node.node_type = 'DUMMY_SCHEMA' # we promote it to a schema for now
  159. mantis_node.inputs.init_sockets(['Schema Length']) # add a Schema Length socket
  160. mantis_node.parameters['Schema Length'] = 1 # set the length to 1 since it is a single group instance
  161. # we'll make the autogenerated nodes for constant inputs. It doesn't matter that there is technically
  162. # a prototype available for each one -- these are cheap and I want this to be easy.
  163. from .readtree import make_connections_to_ng_dummy
  164. make_connections_to_ng_dummy(self.node.base_tree, self.autogen_path_names, frame_mantis_nodes, self.all_nodes, mantis_node)
  165. else:
  166. mantis_node = prototype_mantis_node.__class__(signature, prototype_mantis_node.base_tree, prototype=ui_node)
  167. else:
  168. mantis_node = prototype_mantis_node.__class__(signature, prototype_mantis_node.base_tree)
  169. frame_mantis_nodes[mantis_node.signature] = mantis_node
  170. mantis_node.ui_signature=ui_signature # set the natural signature to ensure we can access from the UI
  171. if mantis_node.prepared == False:
  172. unprepared.append(mantis_node)
  173. if mantis_node.__class__.__name__ in custom_props_types:
  174. setup_custom_props_from_np(mantis_node, prototype_ui_node)
  175. mantis_node.fill_parameters(prototype_ui_node)
  176. # be sure to pass on the Mantis Context to them
  177. mantis_node.mContext=mContext
  178. def handle_link_from_index_input(self, index, frame_mantis_nodes, ui_link):
  179. _from_name, to_name = get_link_in_out(ui_link)
  180. to_node = frame_mantis_nodes[ (*self.autogen_path_names, to_name+self.index_str()) ]
  181. if to_node.node_type in ['DUMMY', 'DUMMY_SCHEMA']:
  182. from .utilities import gen_nc_input_for_data
  183. nc_cls = gen_nc_input_for_data(ui_link.from_socket)
  184. if (nc_cls): #HACK
  185. unique_name = "".join([
  186. ui_link.to_socket.node.name+self.index_str(),
  187. ui_link.from_socket.name,
  188. ui_link.from_socket.identifier,
  189. "==>",
  190. ui_link.to_socket.name,
  191. ui_link.to_socket.identifier,
  192. ])
  193. sig = ("MANTIS_AUTOGENERATED", *self.tree_path_names[1:-1], unique_name)
  194. nc_from = frame_mantis_nodes.get(sig)
  195. if not nc_from:
  196. nc_from = nc_cls(sig, self.node.base_tree)
  197. # ugly! maybe even a HACK!
  198. nc_from.inputs = {}
  199. from .base_definitions import NodeSocket
  200. nc_from.outputs = {ui_link.from_socket.name:NodeSocket(name = ui_link.from_socket.name, node=nc_from)}
  201. nc_from.parameters = {ui_link.from_socket.name:index}
  202. frame_mantis_nodes[sig]=nc_from
  203. nc_from.execution_prepared=True
  204. from_node = nc_from
  205. self.solved_nodes[sig]=from_node
  206. _connection = from_node.outputs[ui_link.from_socket.name].connect(node=to_node, socket=ui_link.to_socket.identifier)
  207. return
  208. # Since the index is already determined, it is safe to remove the socket and just keep the value.
  209. to_node.parameters[ui_link.to_socket.name] = index
  210. del to_node.inputs[ui_link.to_socket.name]
  211. def handle_link_from_schema_length_input(self, frame_mantis_nodes, ui_link):
  212. # see, here I can just use the schema node
  213. _from_name, to_name = get_link_in_out(ui_link)
  214. if to_name in replace_types:
  215. to_node = self.schema_nodes[(*self.autogen_path_names, to_name)]
  216. else:
  217. to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())]
  218. # this self.index_link is only used here?
  219. if self.index_link is None:
  220. # this should be impossible because the Schema gets an auto-generated Int input.
  221. raise NotImplementedError("This code should be unreachable. Please report this as a bug!")
  222. if (self.index_link.from_node):
  223. connection = self.index_link.from_node.outputs[self.index_link.from_socket].connect(node=to_node, socket=ui_link.to_socket.name)
  224. # otherwise we can autogen an int input I guess...?
  225. else:
  226. raise RuntimeError("I was expecting there to be an incoming connection here for Schema Length")
  227. def handle_link_from_incoming_connection_input(self, frame_mantis_nodes, ui_link):
  228. incoming = self.incoming_connections[ui_link.from_socket.name]
  229. if incoming is not None:
  230. from_node = incoming.from_node
  231. _from_name, to_name = get_link_in_out(ui_link)
  232. to_node = frame_mantis_nodes[ (*self.autogen_path_names, to_name+self.index_str()) ]
  233. socket_name=ui_link.to_socket.name
  234. if to_node.node_type in [ 'DUMMY_SCHEMA' ]:
  235. socket_name=ui_link.to_socket.identifier
  236. connection = from_node.outputs[incoming.from_socket].connect(node=to_node, socket=socket_name)
  237. init_connections(from_node)
  238. def handle_link_to_outgoing_connection_output(self, frame_mantis_nodes, ui_link,):
  239. mantis_incoming_node = self.schema_nodes[*self.tree_path_names, 'SchemaIncomingConnection']
  240. for mantis_link in mantis_incoming_node.outputs[ui_link.to_socket.name].links:
  241. to_mantis_node, to_mantis_socket = mantis_link.to_node, mantis_link.to_socket
  242. from_name = get_link_in_out(ui_link)[0]
  243. from_mantis_node = self.solved_nodes[ (*self.autogen_path_names, from_name+self.prev_index_str()) ]
  244. to_mantis_node = frame_mantis_nodes[ (*self.autogen_path_names, to_mantis_node.signature[-1]+self.index_str()) ]
  245. from_socket_name = ui_link.from_socket.name
  246. if from_mantis_node.node_type in ['DUMMY_SCHEMA']:
  247. from_socket_name = ui_link.from_socket.identifier
  248. connection = from_mantis_node.outputs[from_socket_name].connect(node=to_mantis_node, socket=to_mantis_socket)
  249. # We want to delete the links from the tree into the schema node.
  250. # TODO: this is not robust enough and I do not feel sure this is doing the right thing.
  251. if existing_link := self.incoming_connections[ui_link.to_socket.name]:
  252. if existing_link.to_node == self.node:
  253. print ("INFO: Deleting...", existing_link)
  254. if self.node.signature[-1] in existing_link.to_node.signature:
  255. existing_link.die()
  256. # BUG may exist here.
  257. self.incoming_connections[ui_link.to_socket.name] = connection
  258. def handle_link_from_constant_input(self, frame_mantis_nodes, ui_link, to_ui_node):
  259. incoming = self.constant_in[ui_link.from_socket.name]
  260. from_node = incoming.from_node
  261. to_name = get_link_in_out(ui_link)[1]
  262. to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())]
  263. to_socket=ui_link.to_socket.name
  264. from .base_definitions import MantisNodeGroup, SchemaGroup
  265. if isinstance(to_ui_node, (SchemaGroup, MantisNodeGroup)):
  266. to_socket=ui_link.to_socket.identifier
  267. connection = from_node.outputs[incoming.from_socket].connect(node=to_node, socket=to_socket)
  268. init_connections(from_node)
  269. def handle_link_to_array_input_get(self, frame_mantis_nodes, ui_link, index):
  270. from_name, to_name = get_link_in_out(ui_link)
  271. from_nc = frame_mantis_nodes[(*self.autogen_path_names, from_name+self.index_str())]
  272. to_nc = self.schema_nodes[(*self.tree_path_names, to_name)]
  273. # this only needs to be done once:
  274. if index == 0: # BUG? HACK? TODO find out what is going on here.
  275. # Kill the link between the schema node group and the node connecting to it
  276. old_nc = self.all_nodes[(*self.tree_path_names, from_name)]
  277. # I am not sure about this!
  278. existing_link = old_nc.outputs[ui_link.from_socket.name].links[0]
  279. existing_link.die()
  280. #
  281. connection = from_nc.outputs[ui_link.from_socket.name].connect(node=to_nc, socket=ui_link.to_socket.name)
  282. def handle_link_from_array_input(self, frame_mantis_nodes, ui_link, index):
  283. get_index = index
  284. try:
  285. incoming = self.array_input_connections[ui_link.from_socket.identifier][get_index]
  286. except IndexError:
  287. if len(self.array_input_connections[ui_link.from_socket.identifier]) > 0:
  288. incoming = self.array_input_connections[ui_link.from_socket.identifier][0]
  289. # prOrange(incoming.from_node.node_type)
  290. if incoming.from_node.node_type not in ['DUMMY_SCHEMA']:
  291. raise NotImplementedError(wrapRed("dev: make it so Mantis checks if there are enough Array inputs."))
  292. else: # do nothing
  293. return
  294. else:
  295. raise RuntimeError(wrapRed("make it so Mantis checks if there are enough Array inputs!"))
  296. to_name = get_link_in_out(ui_link)[1]
  297. to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())]
  298. connection = incoming.from_node.outputs[incoming.from_socket].connect(node=to_node, socket=ui_link.to_socket.name)
  299. init_connections(incoming.from_node)
  300. def handle_link_from_array_input_all(self, frame_mantis_nodes, ui_link):
  301. all_links = self.array_input_connections[ui_link.from_socket.identifier]
  302. to_name = get_link_in_out(ui_link)[1]
  303. to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())]
  304. # connection = incoming.from_node.outputs[incoming.from_socket].connect(node=to_node, socket=ui_link.to_socket.name)
  305. for l in all_links:
  306. # we need to copy the link with the new from-node info
  307. from .base_definitions import NodeLink
  308. to_socket_name=ui_link.to_socket.name
  309. if to_node.node_type in ['DUMMY_SCHEMA']:
  310. to_socket_name=ui_link.to_socket.identifier
  311. connection = NodeLink(l.from_node, l.from_socket, to_node, to_socket_name, l.multi_input_sort_id)
  312. to_node.flush_links()
  313. def handle_link_to_constant_output(self, frame_mantis_nodes, index, ui_link, to_ui_node):
  314. to_node = self.schema_nodes[(*self.tree_path_names, to_ui_node.bl_idname)]
  315. expose_when = to_node.evaluate_input('Expose when N==')
  316. # HACK here to force it to work with ordinary node groups, which don't seem to set this value correctly.
  317. if to_ui_node.bl_idname == "NodeGroupOutput":
  318. expose_when = index # just set it directly since it is getting set to None somewhere (I should find out where tho)
  319. # end HACK
  320. if index == expose_when:
  321. for outgoing in self.constant_out[ui_link.to_socket.name]:
  322. to_node = outgoing.to_node
  323. from_name = get_link_in_out(ui_link)[0]
  324. from_node = frame_mantis_nodes[(*self.autogen_path_names, from_name+self.index_str()) ]
  325. connection = from_node.outputs[ui_link.from_socket.name].connect(node=to_node, socket=outgoing.to_socket)
  326. # WTF is even happening here?? TODO BUG HACK
  327. def handle_link_to_array_output(self, frame_mantis_nodes, index, ui_link, to_ui_node, from_ui_node):# if this duplicated code works, dedupe!
  328. to_node = self.schema_nodes[(*self.tree_path_names, to_ui_node.bl_idname)] # get it by [], we want a KeyError if this fails
  329. for outgoing in self.array_output_connections[ui_link.to_socket.identifier]:
  330. # print (type(outgoing))
  331. from .schema_containers import SchemaIndex
  332. from_name = get_link_in_out(ui_link)[0]
  333. from_node = frame_mantis_nodes[ (*self.autogen_path_names, from_name+self.index_str()) ]
  334. from_socket_name = ui_link.from_socket.name
  335. if not from_node:
  336. from_node = self.schema_nodes[(*self.tree_path_names, from_ui_node.bl_idname)]
  337. to_node = outgoing.to_node
  338. if isinstance(from_node, SchemaIndex): # I think I need to dedup this stuff
  339. # print("INDEX")
  340. from .utilities import gen_nc_input_for_data
  341. nc_cls = gen_nc_input_for_data(ui_link.from_socket)
  342. if (nc_cls): #HACK
  343. sig = ("MANTIS_AUTOGENERATED", *self.tree_path_names[1:-1], self.index_str(), ui_link.from_socket.name, ui_link.from_socket.identifier)
  344. nc_from = nc_cls(sig, self.node.base_tree)
  345. nc_from.execution_prepared=True
  346. # ugly! maybe even a HACK!
  347. nc_from.inputs = {}
  348. from .node_container_common import NodeSocket
  349. nc_from.outputs = {ui_link.from_socket.name:NodeSocket(name = ui_link.from_socket.name, node=nc_from)}
  350. from .node_container_common import get_socket_value
  351. if ui_link.from_socket.name in ['Index']:
  352. nc_from.parameters = {ui_link.from_socket.name:index}
  353. else:
  354. nc_from.parameters = {ui_link.from_socket.name:self.solve_length}
  355. frame_mantis_nodes[sig]=nc_from
  356. from_node = nc_from
  357. self.solved_nodes[sig]=from_node
  358. elif from_node.node_type == 'DUMMY_SCHEMA':
  359. from_socket_name = ui_link.from_socket.identifier
  360. # I have a feeling that something bad will happen if both of these conditions (above and below) are true
  361. if to_node.node_type == 'DUMMY_SCHEMA' and to_node.prepared:
  362. other_stem = ('SCHEMA_AUTOGENERATED', *to_node.signature[1:])
  363. from .utilities import get_node_prototype
  364. other_schema_np = get_node_prototype(to_node.signature, to_node.base_tree)
  365. other_schema_tree = other_schema_np.node_tree
  366. for n in other_schema_tree.nodes:
  367. if n.bl_idname not in ["SchemaArrayInput", "SchemaArrayInputGet"]:
  368. continue
  369. out = n.outputs[outgoing.to_socket]
  370. for l in out.links:
  371. other_index_str = lambda : '.'+str(to_node.uuid)+'.'+str(index).zfill(4)
  372. # get it by [], we want a KeyError if this fails
  373. try:
  374. out_node = self.all_nodes[(*other_stem, l.to_node.name+other_index_str())]
  375. except KeyError as e:
  376. for n in self.all_nodes:
  377. if len(n) > len(other_stem)+1: break
  378. for elem in other_stem:
  379. if elem not in n: break
  380. else:
  381. print(n)
  382. raise e
  383. connection = from_node.outputs[from_socket_name].connect(node=out_node, socket=l.to_socket.name)
  384. else:
  385. connection = from_node.outputs[from_socket_name].connect(node=to_node, socket=outgoing.to_socket)
  386. def handle_link_from_array_input_get(self, frame_mantis_nodes, index, ui_link, from_ui_node ):
  387. get_index = index
  388. from_node = self.schema_nodes[(*self.tree_path_names, from_ui_node.bl_idname)]
  389. from .utilities import cap, wrap
  390. get_index = from_node.evaluate_input("Index", index)
  391. oob = from_node.evaluate_input("OoB Behaviour")
  392. # we must assume that the array has sent the correct number of links
  393. if oob == 'WRAP':
  394. get_index = wrap(get_index, len(self.array_input_connections[ui_link.from_socket.identifier])-1, 0)
  395. if oob == 'HOLD':
  396. get_index = cap(get_index, len(self.array_input_connections[ui_link.from_socket.identifier])-1)
  397. try:
  398. incoming = self.array_input_connections[ui_link.from_socket.identifier][get_index]
  399. except IndexError:
  400. raise RuntimeError(wrapRed("Dummy! You need to make it so Mantis checks if there are enough Array inputs! It should probably have a Get Index!"))
  401. to_name = get_link_in_out(ui_link)[1]
  402. to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())]
  403. connection = incoming.from_node.outputs[incoming.from_socket].connect(node=to_node, socket=ui_link.to_socket.name)
  404. init_connections(incoming.from_node)
  405. def test_is_sub_schema(self, other):
  406. for i in range(len(other.ui_signature)-1): # -1, we don't want to check this node, obviously
  407. if self.node.ui_signature[:i+1]:
  408. return False
  409. return True
  410. def prepare_nodes(self, unprepared):
  411. # At this point, we've already run a pretty exhaustive preperation phase to prep the schema's dependencies
  412. # So we should not need to add any new dependencies unless there is a bug elsewhere.
  413. # and in fact, I could skip this in some cases, and should investigate if profiling reveals a slowdown here.
  414. forbidden=set()
  415. # forbid some nodes - they aren't necessary to solve the schema & cause problems.
  416. from .readtree import execution_error_cleanup
  417. while unprepared:
  418. nc = unprepared.pop()
  419. if sum([dep.prepared for dep in nc.hierarchy_dependencies]) == len(nc.hierarchy_dependencies):
  420. try:
  421. nc.bPrepare()
  422. except Exception as e:
  423. execution_error_cleanup(nc, e)
  424. if nc.node_type == 'DUMMY_SCHEMA':
  425. self.solve_nested_schema(nc)
  426. elif nc.node_type == 'DUMMY_SCHEMA' and not self.test_is_sub_schema(nc):
  427. forbidden.add(nc)
  428. continue # do NOT add this as a dependency.
  429. else: # Keeping this for-loop as a fallback, it should never add dependencies though
  430. can_add_me = True
  431. for dep in nc.hierarchy_dependencies:
  432. if not dep.prepared and dep not in unprepared:
  433. if dep in forbidden:
  434. can_add_me=False
  435. forbidden.add(nc) # forbid the parent, too
  436. continue
  437. prOrange(f"Adding dependency... {dep}")
  438. unprepared.appendleft(dep)
  439. if can_add_me:
  440. unprepared.appendleft(nc) # just rotate them until they are ready.
  441. def solve_iteration(self):
  442. """ Solve an iteration of the schema.
  443. - 1 Create the Mantis Node instances that represent this iteration of the schema
  444. - 2 Connect the links from the entrypoint or previous iteration.
  445. - 3 Connect the constant and array links, and any link between nodes entirely within the tree
  446. - 4 Prepare the nodes that modify data (in case of e.g. array get index or nested schema length input)
  447. - 5 Connect the final prepared nodes
  448. and return the nodes that were created in this schema iteration (frame).
  449. This function also adds to held_links to pass data between iterations.
  450. """
  451. from .schema_definitions import (SchemaIndex,
  452. SchemaArrayInput,
  453. SchemaArrayInputGet,
  454. SchemaArrayInputAll,
  455. SchemaArrayOutput,
  456. SchemaConstInput,
  457. SchemaConstOutput,
  458. SchemaOutgoingConnection,
  459. SchemaIncomingConnection,)
  460. from .utilities import clear_reroutes, link_node_containers
  461. self.set_index_strings()
  462. frame_mantis_nodes = {}
  463. # Later we have to run bPrepare() on these guys, so we make the deque and fill it now.
  464. from collections import deque
  465. unprepared= deque()
  466. self.gen_solve_iteration_mantis_nodes(frame_mantis_nodes, unprepared)
  467. # This is where we handle node connections BETWEEN frames
  468. while(self.held_links):
  469. ui_link = self.held_links.pop()
  470. to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node
  471. if isinstance(to_ui_node, SchemaOutgoingConnection):
  472. self.handle_link_to_outgoing_connection_output(frame_mantis_nodes, ui_link)
  473. # Get the rerouted links from the graph. We don't really need to do this every iteration.
  474. # TODO: use profiling to determine if this is slow; if so: copy & reuse the data, refactor the pop()'s out.
  475. ui_links = clear_reroutes(list(self.tree.links))
  476. # Now we handle ui_links in the current frame, including those ui_links between Schema nodes and "real" nodes
  477. awaiting_prep_stage = []
  478. for ui_link in ui_links:
  479. to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node
  480. if isinstance(from_ui_node, SchemaIndex):
  481. if ui_link.from_socket.name == "Index":
  482. self.handle_link_from_index_input(self.index, frame_mantis_nodes, ui_link)
  483. elif ui_link.from_socket.name == "Schema Length":
  484. self.handle_link_from_schema_length_input(frame_mantis_nodes, ui_link)
  485. continue
  486. if isinstance(from_ui_node, SchemaIncomingConnection):
  487. if ui_link.from_socket.name in self.incoming_connections.keys():
  488. self.handle_link_from_incoming_connection_input(frame_mantis_nodes, ui_link)
  489. continue
  490. if isinstance(from_ui_node, (SchemaConstInput, NodeGroupInput)):
  491. if ui_link.from_socket.name in self.constant_in.keys():
  492. self.handle_link_from_constant_input( frame_mantis_nodes, ui_link, to_ui_node)
  493. continue
  494. if isinstance(to_ui_node, SchemaArrayInputGet):
  495. self.handle_link_to_array_input_get( frame_mantis_nodes, ui_link, self.index)
  496. continue
  497. if isinstance(from_ui_node, SchemaArrayInput):
  498. self.handle_link_from_array_input(frame_mantis_nodes, ui_link, self.index)
  499. continue
  500. if isinstance(from_ui_node, SchemaArrayInputAll):
  501. self.handle_link_from_array_input_all(frame_mantis_nodes, ui_link)
  502. continue
  503. # HOLD these links to the next iteration:
  504. if isinstance(to_ui_node, SchemaOutgoingConnection):
  505. self.held_links.append(ui_link)
  506. continue
  507. # HOLD these links until prep is done a little later
  508. if isinstance(to_ui_node, (SchemaConstOutput, NodeGroupOutput)) or isinstance(to_ui_node, SchemaArrayOutput) or \
  509. isinstance(from_ui_node, SchemaArrayInputGet):
  510. awaiting_prep_stage.append(ui_link)
  511. continue
  512. # for any of the special cases, we hit a 'continue' block. So this connection is not special, and is made here.
  513. connection = link_node_containers(self.autogen_path_names, ui_link, frame_mantis_nodes, from_suffix=self.index_str(), to_suffix=self.index_str())
  514. for signature, node in frame_mantis_nodes.items():
  515. self.solved_nodes[signature]=node
  516. if node.node_type == "DUMMY_SCHEMA":
  517. # make sure to add the nodes to the group's sockets if the user set them directly
  518. from .readtree import make_connections_to_ng_dummy
  519. make_connections_to_ng_dummy(
  520. self.node.base_tree,
  521. self.autogen_path_names,
  522. {}, # just pass an empty dict, this argument is not needed in this context
  523. self.all_nodes,
  524. node)
  525. from .utilities import init_schema_dependencies
  526. init_schema_dependencies(node, self.all_nodes)
  527. else:
  528. init_dependencies(node) # it is hard to overstate how important this single line of code is
  529. self.prepare_nodes(unprepared)
  530. while(awaiting_prep_stage):
  531. ui_link = awaiting_prep_stage.pop()
  532. to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node
  533. if isinstance(to_ui_node, (SchemaConstOutput, NodeGroupOutput)):
  534. self.handle_link_to_constant_output(frame_mantis_nodes, self.index, ui_link, to_ui_node)
  535. if isinstance(to_ui_node, SchemaArrayOutput):
  536. self.handle_link_to_array_output(frame_mantis_nodes, self.index, ui_link, to_ui_node, from_ui_node)
  537. if isinstance(from_ui_node, SchemaArrayInputGet):
  538. self.handle_link_from_array_input_get(frame_mantis_nodes, self.index, ui_link, from_ui_node )
  539. # end seciton
  540. return frame_mantis_nodes
  541. def solve_nested_schema(self, schema_nc):
  542. """ Solves all schema node groups found in this Schema. This is a recursive function, which will
  543. solve all levels of nested schema - since this function is called by solver.solve().
  544. """
  545. solver=None
  546. if schema_nc.prepared == False:
  547. all_nodes = self.all_nodes.copy()
  548. ui_node = schema_nc.prototype
  549. length = schema_nc.evaluate_input("Schema Length")
  550. tree = ui_node.node_tree
  551. if schema_nc.prototype.bl_idname == "MantisNodeGroup":
  552. prOrange(f"Expanding Node Group {tree.name} in node {schema_nc}.")
  553. else:
  554. prOrange(f"Expanding schema {tree.name} in node {schema_nc} with length {length}.")
  555. solver = SchemaSolver(schema_nc, all_nodes, ui_node, schema_nc.ui_signature)
  556. solved_nodes = solver.solve()
  557. schema_nc.prepared = True
  558. for k,v in solved_nodes.items():
  559. self.solved_nodes[k]=v
  560. return solver
  561. def finalize(self, frame_nc):
  562. from .schema_definitions import (SchemaOutgoingConnection,)
  563. for i in range(len(self.held_links)):
  564. link = self.held_links.pop()
  565. to_np = link.to_socket.node; from_np = link.from_socket.node
  566. if isinstance(to_np, SchemaOutgoingConnection):
  567. if link.to_socket.name in self.outgoing_connections.keys():
  568. if (outgoing_links := self.outgoing_connections[link.to_socket.name]) is None: continue
  569. for outgoing in outgoing_links:
  570. if outgoing:
  571. to_node = outgoing.to_node
  572. from_node =frame_nc[(*self.autogen_path_names, from_np.name+self.index_str()) ]
  573. from_socket_name = link.from_socket.name
  574. if from_node.node_type in ['DUMMY_SCHEMA']:
  575. from_socket_name = link.from_socket.identifier
  576. connection = from_node.outputs[from_socket_name].connect(node=to_node, socket=outgoing.to_socket)
  577. # we need to kill the link between the Schema itself and the next node and update the deps. Otherwise:confusing bugs.
  578. outgoing.die(); init_dependencies(to_node)
  579. # else: # the node just isn't connected out this socket.
  580. # # solve all unsolved nested schemas...
  581. for schema_sig, schema_nc in self.nested_schemas.items():
  582. self.solve_nested_schema(schema_nc)
  583. for n in self.autogenerated_nodes.values():
  584. init_connections(n)
  585. for c in n.connections:
  586. init_dependencies(c)
  587. all_outgoing_links = []
  588. for conn in self.outgoing_connections.values():
  589. for outgoing in conn:
  590. all_outgoing_links.append(outgoing)
  591. for conn in self.constant_out.values():
  592. for outgoing in conn:
  593. all_outgoing_links.append(outgoing)
  594. for conn in self.array_output_connections.values():
  595. for outgoing in conn:
  596. all_outgoing_links.append(outgoing)
  597. for outgoing in all_outgoing_links:
  598. to_node = outgoing.to_node
  599. for l in to_node.inputs[outgoing.to_socket].links:
  600. if self.node == l.from_node:
  601. l.die()
  602. for inp in self.node.inputs.values():
  603. for l in inp.links:
  604. init_connections(l.from_node) # to force it to have hierarchy connections with the new nodes.
  605. def solve(self):
  606. if self.solve_length < 1:
  607. from .base_definitions import GraphError
  608. for o in self.node.outputs:
  609. if o.is_linked:
  610. raise GraphError(f"ERROR: Schema {self} has a length"
  611. " of 0 but other nodes depend on it.")
  612. print (f"WARN: Schema {self} has a length of 0 or less and will not expand.")
  613. return {} # just don't do anything - it's OK to have a noop schema if it doesn't have dependencies.
  614. for index in range(self.solve_length):
  615. self.index = index
  616. frame_mantis_nodes = self.solve_iteration()
  617. for sig, nc in frame_mantis_nodes.items():
  618. if nc.node_type == 'DUMMY_SCHEMA':
  619. self.nested_schemas[sig] = nc
  620. self.finalize(frame_mantis_nodes)
  621. self.node.solver = self
  622. self.node.prepared = True
  623. return self.solved_nodes