schema_solve.py 50 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857
  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, \
  7. MantisNodeGroup, SchemaGroup, replace_types, GraphError, links_sort_key)
  8. from .node_container_common import setup_custom_props_from_np
  9. # a class that solves Schema nodes
  10. from bpy.types import NodeGroupInput, NodeGroupOutput
  11. from .readtree import execution_error_cleanup
  12. class SchemaSolver:
  13. def __init__(self, schema_dummy, nodes, prototype, signature=None, error_popups=False):
  14. self.all_nodes = nodes # this is the parsed tree from Mantis
  15. self.node = schema_dummy
  16. self.node.solver = self
  17. self.solved = False
  18. self.tree = prototype.node_tree
  19. self.uuid = self.node.uuid
  20. self.error_popups = error_popups
  21. if signature:
  22. self.signature = signature
  23. else:
  24. self.signature = self.node.signature
  25. self.schema_nodes={}
  26. self.solved_nodes = {}
  27. self.incoming_connections = {}
  28. self.outgoing_connections = {}
  29. self.constant_in = {}
  30. self.constant_out = {}
  31. self.array_input_connections = {}
  32. self.array_output_connections = {}
  33. self.nested_schemas = {}
  34. self.autogenerated_nodes = {}
  35. self.held_links = []
  36. self.tree_path_names = [*self.node.signature] # same tree as the schema node
  37. self.autogen_path_names = ['SCHEMA_AUTOGENERATED', *self.node.signature[1:]]
  38. self.is_node_group = False
  39. if self.node.prototype.bl_idname == "MantisNodeGroup":
  40. self.is_node_group = True
  41. if self.node.inputs['Schema Length'].links:
  42. self.index_link = self.node.inputs['Schema Length'].links[0]
  43. else:
  44. self.index_link = None
  45. # this should never happen, but it's OK to check.
  46. if self.node.inputs["Schema Length"].is_linked:
  47. if (other := self.node.inputs['Schema Length'].links[0].from_node).prepared == False:
  48. raise RuntimeError(f"Schema Length cannot be determined for {self.node} because {other} is not prepared.")
  49. self.solve_length = self.node.evaluate_input("Schema Length")
  50. # I'm making this a property of the solver because the solver's data is modified as it solves each iteration
  51. self.index = 0
  52. prWhite(f"\nExpanding schema {self.tree.name} in node {self.node.signature}"
  53. f" with length {self.solve_length}.")
  54. self.init_schema_links()
  55. self.set_index_strings()
  56. # Sort the multi-input nodes in reverse order of ID, this ensures that they are
  57. # read in the order they were created
  58. for inp in self.node.inputs.values():
  59. inp.links.sort(key=links_sort_key)
  60. from bpy.types import NodeGroupInput, NodeGroupOutput
  61. for ui_node in self.tree.nodes:
  62. # first we need to fill the parameters of the schema nodes.
  63. # we use the bl_idname because all schema nodes should be single-instance
  64. signature = (*self.tree_path_names, ui_node.bl_idname)
  65. if isinstance(ui_node, (SchemaUINode, NodeGroupInput, NodeGroupOutput)):
  66. # We use the schema node's "natural signature" here because it represents
  67. # the "original" signature of the schema UI group node since this schema
  68. # solver may be in a nested schema, and its node's signature may have
  69. # uuid/index attached.
  70. get_sig = (*self.node.ui_signature, ui_node.bl_idname)
  71. if not (mantis_node := self.all_nodes.get(get_sig)):
  72. raise RuntimeError(wrapRed(f"Not found: {get_sig}"))
  73. self.schema_nodes[signature] = mantis_node
  74. mantis_node.fill_parameters(ui_node)
  75. # HACK to make Group Nodes work
  76. if ui_node.bl_idname == "NodeGroupInput":
  77. from .schema_nodes import SchemaConstInput
  78. mantis_node = SchemaConstInput(signature=signature, base_tree=self.node.base_tree, parent_schema_node=self.node)
  79. self.schema_nodes[signature] = mantis_node
  80. mantis_node.fill_parameters(ui_node)
  81. if ui_node.bl_idname == "NodeGroupOutput":
  82. from .schema_nodes import SchemaConstOutput
  83. mantis_node = SchemaConstOutput(signature=signature, base_tree=self.node.base_tree, parent_schema_node=self.node)
  84. self.schema_nodes[signature] = mantis_node
  85. mantis_node.fill_parameters(ui_node)
  86. interface = self.setup_interface_node(signature=( *self.signature, "InputInterface"))
  87. def setup_interface_node(self, signature):
  88. from .internal_containers import GroupInterface
  89. interface = GroupInterface(
  90. signature,
  91. self.node.base_tree, self.node.prototype, 'INPUT',)
  92. # now we need to connect it up
  93. # adapted from readtree.grp_node_reroute_common
  94. from .base_definitions import links_sort_key
  95. for in_node_input in self.node.inputs:
  96. if in_node_input.name != 'Schema Length':
  97. for interface_socket in self.tree.interface.items_tree:
  98. if interface_socket.item_type == 'PANEL': continue
  99. if interface_socket.identifier == in_node_input.name: break
  100. if hasattr(interface_socket, "is_array") and interface_socket.is_array: continue
  101. if hasattr(interface_socket, "is_connection") and interface_socket.is_connection: continue
  102. # for now, do not try and connect to the array/connection inputs.
  103. i = 0
  104. if len(in_node_input.links)>1: # sort here to ensure correct sub_sort_id
  105. in_node_input.links.sort(key=links_sort_key)
  106. for in_link in in_node_input.links:
  107. from_node = in_link.from_node; from_socket = in_link.from_socket
  108. link = from_node.outputs[from_socket].connect(
  109. interface,in_node_input.name, sort_id = 2**16, sub_sort_id=i)
  110. i += 1
  111. in_node_input.links.sort(key=links_sort_key)
  112. return interface # I want to be able to set this up elsewhere
  113. def set_index_strings(self):
  114. self.index_str = lambda : '.'+str(self.uuid)+'.'+str(self.index).zfill(4)
  115. self.prev_index_str = lambda : '.'+str(self.uuid)+'.'+str(self.index-1).zfill(4)
  116. if self.is_node_group:
  117. self.index_str=lambda : ''
  118. self.prev_index_str=lambda : ''
  119. def init_schema_links(self,):
  120. """ Sort and store the links to/from the Schema group node."""
  121. for item in self.tree.interface.items_tree:
  122. if item.item_type == 'PANEL': continue
  123. from .utilities import read_schema_type
  124. parent_name = read_schema_type(item)
  125. # just gonna try and make the
  126. match parent_name:
  127. case 'Connection':
  128. if item.in_out == 'INPUT':
  129. if incoming_links := self.node.inputs[item.identifier].links:
  130. self.incoming_connections[item.name] = incoming_links[0]
  131. else:
  132. self.incoming_connections[item.name] = None
  133. else: # OUTPUT
  134. if outgoing_links := self.node.outputs[item.identifier].links:
  135. self.outgoing_connections[item.name] = outgoing_links.copy()
  136. else:
  137. self.outgoing_connections[item.name] = []
  138. case 'Constant':
  139. if item.in_out == 'INPUT':
  140. if constant_in_links := self.node.inputs[item.identifier].links:
  141. self.constant_in[item.name] = constant_in_links[0]
  142. else:
  143. self.constant_in[item.name] = None
  144. else: # OUTPUT
  145. if constant_out_links := self.node.outputs[item.identifier].links:
  146. self.constant_out[item.name] = constant_out_links.copy()
  147. else:
  148. self.constant_out[item.name] = []
  149. case 'Array':
  150. if item.in_out == 'INPUT':
  151. if item.identifier not in self.array_input_connections.keys():
  152. self.array_input_connections[item.identifier]=[]
  153. if in_links := self.node.inputs[item.identifier].links:
  154. self.array_input_connections[item.identifier]=in_links.copy()
  155. # I am tempted to put a check here, but it is sufficient to
  156. # rely on hierarchy links ensuring the arrays are prepared.
  157. # and testing here for un-prepared arrays will mess up schema
  158. # relationships in non-hierarchy situations.
  159. # Still, I'd like to have an easier time catching these problems.
  160. else: # OUTPUT
  161. if item.identifier not in self.array_output_connections.keys():
  162. self.array_output_connections[item.identifier]=[]
  163. if out_links := self.node.outputs[item.identifier].links:
  164. self.array_output_connections[item.identifier] = out_links.copy()
  165. def is_node_deeper_nested(self, queried_node, compare_node):
  166. return len(compare_node.signature) < len(queried_node.signature)
  167. def gen_solve_iteration_mantis_nodes(self, frame_mantis_nodes, unprepared):
  168. for prototype_ui_node in self.tree.nodes:
  169. mantis_node_name = prototype_ui_node.name
  170. index_str = self.index_str()
  171. mContext=self.node.mContext
  172. if isinstance(prototype_ui_node, SchemaUINode):
  173. continue # IGNORE the schema interface nodes, we already made them in __init__()
  174. # they are reused for each iteration.
  175. elif prototype_ui_node.bl_idname in ['NodeFrame', 'NodeReroute']:
  176. continue # IGNORE stuff that is purely UI - frames, reroutes.
  177. elif prototype_ui_node.bl_idname in ['NodeGroupInput', 'NodeGroupOutput']:
  178. continue # we converted these to Schema Nodes because they represent a Group input.
  179. signature = (*self.autogen_path_names, mantis_node_name+index_str)
  180. ui_signature=(*self.signature, mantis_node_name)
  181. prototype_mantis_node = self.all_nodes[ui_signature]
  182. # the prototype_mantis_node was generated inside the schema when we parsed the tree.
  183. # it is the prototype of the mantis node which we make for this iteration
  184. # for Schema sub-nodes ... they need a prototype to init.
  185. if prototype_mantis_node.node_type in ['DUMMY', 'DUMMY_SCHEMA']:
  186. # We stored the prototype ui_node when creating the Mantis node.
  187. ui_node = prototype_mantis_node.prototype
  188. # if prototype_mantis_node is a group or schema: TODO changes are needed elsewhere to make this easier to read. LEGIBILITY
  189. if ui_node.bl_idname in ["MantisNodeGroup", "MantisSchemaGroup"]:
  190. mantis_node = prototype_mantis_node.__class__(
  191. signature, prototype_mantis_node.base_tree, prototype=ui_node,
  192. ui_signature = prototype_mantis_node.signature)
  193. # now let's copy the links from the prototype node
  194. if ui_node.bl_idname in ["MantisNodeGroup"]:
  195. mantis_node.prepared = False
  196. mantis_node.node_type = 'DUMMY_SCHEMA' # we promote it to a schema for now
  197. mantis_node.inputs.init_sockets(['Schema Length']) # add a Schema Length socket
  198. mantis_node.parameters['Schema Length'] = 1 # set the length to 1 since it is a single group instance
  199. # we'll make the autogenerated nodes for constant inputs. It doesn't matter that there is technically
  200. # a prototype available for each one -- these are cheap and I want this to be easy.
  201. from .readtree import make_connections_to_ng_dummy
  202. make_connections_to_ng_dummy(self.node.base_tree, self.autogen_path_names, frame_mantis_nodes, self.all_nodes, mantis_node)
  203. else:
  204. mantis_node = prototype_mantis_node.__class__(signature, prototype_mantis_node.base_tree, prototype=ui_node)
  205. else:
  206. mantis_node = prototype_mantis_node.__class__(signature, prototype_mantis_node.base_tree)
  207. frame_mantis_nodes[mantis_node.signature] = mantis_node
  208. self.all_nodes[mantis_node.signature] = mantis_node
  209. mantis_node.ui_signature=ui_signature # set the natural signature to ensure we can access from the UI
  210. if mantis_node.prepared == False:
  211. unprepared.append(mantis_node)
  212. if mantis_node.__class__.__name__ in custom_props_types:
  213. setup_custom_props_from_np(mantis_node, prototype_ui_node)
  214. mantis_node.fill_parameters(prototype_ui_node)
  215. # be sure to pass on the Mantis Context to them
  216. mantis_node.mContext=mContext
  217. # set up an interface node at this point.
  218. interface = self.setup_interface_node(signature=( *self.signature, "InputInterface"+self.index_str() ))
  219. interface.mContext = mContext
  220. frame_mantis_nodes[interface.signature] = interface
  221. self.all_nodes[interface.signature] = interface
  222. def handle_link_from_index_input(self, index, frame_mantis_nodes, ui_link):
  223. from .base_definitions import can_remove_socket_for_autogen
  224. _from_name, to_name = get_link_in_out(ui_link)
  225. to_node = frame_mantis_nodes[ (*self.autogen_path_names, to_name+self.index_str()) ]
  226. if (not can_remove_socket_for_autogen(to_node, ui_link.to_socket.name)) or \
  227. to_node.node_type in ['DUMMY', 'DUMMY_SCHEMA']:
  228. from .readtree import autogen_node
  229. unique_name = "".join([
  230. ui_link.to_socket.node.name+self.index_str(),
  231. ui_link.from_socket.name, ui_link.from_socket.identifier,
  232. "==>", ui_link.to_socket.name, ui_link.to_socket.identifier,],)
  233. signature = ("MANTIS_AUTOGENERATED", *self.tree_path_names[1:-1], unique_name)
  234. from_node = self.all_nodes.get(signature)
  235. if not from_node:
  236. from_node = autogen_node(self.node.base_tree, ui_link.from_socket,
  237. signature=signature, mContext=self.node.mContext)
  238. from_node.parameters = {ui_link.from_socket.name:index}
  239. frame_mantis_nodes[signature]=from_node; self.solved_nodes[signature]=from_node
  240. self.all_nodes[signature]=from_node
  241. _connection = from_node.outputs[ui_link.from_socket.name].connect(node=to_node, socket=ui_link.to_socket.identifier)
  242. return
  243. # Since the index is already determined, it is safe to remove the socket and just keep the value.
  244. to_node.parameters[ui_link.to_socket.name] = index
  245. del to_node.inputs[ui_link.to_socket.name]
  246. def handle_link_from_schema_length_input(self, frame_mantis_nodes, ui_link):
  247. # see, here I can just use the schema node
  248. _from_name, to_name = get_link_in_out(ui_link)
  249. if to_name in replace_types:
  250. to_node = self.schema_nodes[(*self.autogen_path_names, to_name)]
  251. else:
  252. to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())]
  253. # this self.index_link is only used here?
  254. if self.index_link is None:
  255. # this should be impossible because the Schema gets an auto-generated Int input.
  256. raise NotImplementedError(" 241 This code should be unreachable. Please report this as a bug!")
  257. if (self.index_link.from_node):
  258. connection = self.index_link.from_node.outputs[self.index_link.from_socket].connect(node=to_node, socket=ui_link.to_socket.name)
  259. # otherwise we can autogen an int input I guess...?
  260. else:
  261. raise RuntimeError("247 I This code should be unreachable. Please report this as a bug!")
  262. # TODO: link the incoming connection to the current Interface node
  263. # It is complicated because the Interface needs the identifier
  264. # but most of this stuff uses the name and I donno how to get the ID
  265. # this isn't very important. I'll add this feature when someone tries
  266. # to use it and complains.
  267. def handle_link_from_incoming_connection_input(self, frame_mantis_nodes, ui_link):
  268. incoming = self.incoming_connections[ui_link.from_socket.name]
  269. if incoming is not None:
  270. from_node = incoming.from_node
  271. _from_name, to_name = get_link_in_out(ui_link)
  272. to_node = frame_mantis_nodes[ (*self.autogen_path_names, to_name+self.index_str()) ]
  273. socket_name=ui_link.to_socket.name
  274. if to_node.node_type in [ 'DUMMY_SCHEMA' ]:
  275. socket_name=ui_link.to_socket.identifier
  276. connection = from_node.outputs[incoming.from_socket].connect(node=to_node, socket=socket_name)
  277. init_connections(from_node)
  278. def handle_link_to_outgoing_connection_output(self, frame_mantis_nodes, ui_link,):
  279. mantis_incoming_node = self.schema_nodes[*self.tree_path_names, 'SchemaIncomingConnection']
  280. for mantis_link in mantis_incoming_node.outputs[ui_link.to_socket.name].links:
  281. to_mantis_node, to_socket_name = mantis_link.to_node, mantis_link.to_socket
  282. from_name = get_link_in_out(ui_link)[0]
  283. from_mantis_node = self.solved_nodes[ (*self.autogen_path_names, from_name+self.prev_index_str()) ]
  284. to_mantis_node_signature = ( *self.autogen_path_names,
  285. to_mantis_node.signature[-1] + self.index_str() )
  286. # we need to detect if the next node is in a group.
  287. # REMEMBER: at this point, nested groups haven't been solved yet.
  288. if to_mantis_node_signature not in frame_mantis_nodes.keys() and \
  289. self.is_node_deeper_nested(to_mantis_node, from_mantis_node):
  290. to_mantis_node = frame_mantis_nodes[ (
  291. *self.autogen_path_names, # the SCHEMA_AUTOGENERATED string
  292. to_mantis_node.signature[-2]+ self.index_str() ) ]
  293. # this connection was forbidden before, right? so this should be a safe assumption.
  294. assert to_mantis_node.node_type == 'DUMMY_SCHEMA', "I expected this to be a group/schema"
  295. # we need to get the identifier of the named socket now.
  296. to_socket_name = to_mantis_node.prototype.inputs[to_socket_name].identifier
  297. else:
  298. to_mantis_node = frame_mantis_nodes[ to_mantis_node_signature ] # add the index string
  299. from_socket_name = ui_link.from_socket.name
  300. if from_mantis_node.node_type in ['DUMMY_SCHEMA']:
  301. from_socket_name = ui_link.from_socket.identifier
  302. connection = from_mantis_node.outputs[from_socket_name].connect(node=to_mantis_node, socket=to_socket_name)
  303. # We want to delete the links from the tree into the schema node.
  304. # TODO: this is not robust enough and I do not feel sure this is doing the right thing.
  305. if existing_link := self.incoming_connections[ui_link.to_socket.name]:
  306. if existing_link.to_node == self.node:
  307. print ("INFO: Deleting...", existing_link)
  308. if self.node.signature[-1] in existing_link.to_node.signature:
  309. existing_link.die()
  310. # BUG may exist here.
  311. self.incoming_connections[ui_link.to_socket.name] = connection
  312. def handle_link_from_constant_input(self, frame_mantis_nodes, ui_link, to_ui_node):
  313. incoming = self.constant_in[ui_link.from_socket.name]
  314. if incoming is None:
  315. raise GraphError (f"{self.node} is missing a required input: {ui_link.from_socket.name}")
  316. from_node = incoming.from_node
  317. to_name = get_link_in_out(ui_link)[1]
  318. to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())]
  319. to_socket=ui_link.to_socket.name
  320. from .base_definitions import MantisNodeGroup, SchemaGroup
  321. if isinstance(to_ui_node, (SchemaGroup, MantisNodeGroup)):
  322. to_socket=ui_link.to_socket.identifier
  323. connection = from_node.outputs[incoming.from_socket].connect(node=to_node, socket=to_socket)
  324. init_connections(from_node)
  325. def handle_link_from_array_input_get(self, frame_mantis_nodes, ui_link ):
  326. from_ui_node = ui_link.from_socket.node
  327. from_node = self.schema_nodes[(*self.node.ui_signature, from_ui_node.bl_idname)]
  328. from collections import deque
  329. unprepared = deque(from_node.hierarchy_dependencies)
  330. self.prepare_nodes(unprepared)
  331. from .utilities import cap, wrap
  332. get_index = from_node.evaluate_input("Index", self.index) # get the most recent link
  333. # getting the link at self.index just saves the trouble of killing the old links
  334. # that are left because this node is reused in each iteration of the schema
  335. assert(get_index is not None), f"Cannot get index in {from_node}"
  336. oob = from_node.evaluate_input("OoB Behaviour")
  337. array_length = len(self.array_input_connections[ui_link.from_socket.identifier])-1
  338. if oob == 'WRAP':
  339. get_index = wrap(0, array_length+1, get_index)
  340. if oob == 'HOLD':
  341. get_index = cap(get_index, array_length)
  342. try:
  343. incoming = self.array_input_connections[ui_link.from_socket.identifier][get_index]
  344. except IndexError:
  345. raise RuntimeError(f"The array index {get_index} is out of bounds. Size: "
  346. f"{len(self.array_input_connections[ui_link.from_socket.identifier])}")
  347. to_name = get_link_in_out(ui_link)[1]
  348. to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())]
  349. from_socket = incoming.from_node.outputs[incoming.from_socket]
  350. connection = from_socket.connect(to_node, ui_link.to_socket.name)
  351. # TODO: kill links to the array get that are no longer needed
  352. # right now I inefficiently waste cpu-time solving extra nodes
  353. # because I took the lazy way out
  354. def handle_link_to_array_input_get(self, frame_mantis_nodes, ui_link):
  355. from_name, to_name = get_link_in_out(ui_link)
  356. from_nc = frame_mantis_nodes[(*self.autogen_path_names, from_name+self.index_str())]
  357. to_nc = self.schema_nodes[(*self.tree_path_names, to_name)]
  358. # this only needs to be done once:
  359. if self.index == 0: # BUG? HACK? TODO find out what is going on here.
  360. # Kill the link between the schema node group and the node connecting to it
  361. old_nc = self.all_nodes[(*self.tree_path_names, from_name)]
  362. # I am not sure about this!
  363. existing_link = old_nc.outputs[ui_link.from_socket.name].links[0]
  364. existing_link.die()
  365. #
  366. connection = from_nc.outputs[ui_link.from_socket.name].connect(node=to_nc, socket=ui_link.to_socket.name)
  367. connection.is_hierarchy = True # we have to mark this because links to Schema are not usually hierarchy (?)
  368. to_nc.hierarchy_dependencies.append(from_nc); from_nc.hierarchy_connections.append(to_nc)
  369. # TODO: review the wisdom of this default.
  370. def handle_link_from_array_input(self, frame_mantis_nodes, ui_link, index):
  371. get_index = index
  372. try:
  373. incoming = self.array_input_connections[ui_link.from_socket.identifier][get_index]
  374. except IndexError:
  375. if len(self.array_input_connections[ui_link.from_socket.identifier]) > 0:
  376. incoming = self.array_input_connections[ui_link.from_socket.identifier][0]
  377. # prOrange(incoming.from_node.node_type)
  378. if incoming.from_node.node_type not in ['DUMMY_SCHEMA']:
  379. raise NotImplementedError(wrapRed("dev: make it so Mantis checks if there are enough Array inputs."))
  380. else: # do nothing
  381. return
  382. else:
  383. raise RuntimeError(wrapRed("make it so Mantis checks if there are enough Array inputs!"))
  384. to_name = get_link_in_out(ui_link)[1]
  385. to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())]
  386. connection = incoming.from_node.outputs[incoming.from_socket].connect(node=to_node, socket=ui_link.to_socket.name)
  387. init_connections(incoming.from_node)
  388. def handle_link_from_array_input_all(self, frame_mantis_nodes, ui_link):
  389. all_links = self.array_input_connections[ui_link.from_socket.identifier]
  390. to_name = get_link_in_out(ui_link)[1]
  391. to_node = frame_mantis_nodes[(*self.autogen_path_names, to_name+self.index_str())]
  392. # connection = incoming.from_node.outputs[incoming.from_socket].connect(node=to_node, socket=ui_link.to_socket.name)
  393. for l in all_links:
  394. # we need to copy the link with the new from-node info
  395. from .base_definitions import NodeLink
  396. to_socket_name=ui_link.to_socket.name
  397. if to_node.node_type in ['DUMMY_SCHEMA']:
  398. to_socket_name=ui_link.to_socket.identifier
  399. connection = NodeLink(l.from_node, l.from_socket, to_node, to_socket_name,
  400. l.multi_input_sort_id, l.sub_sort_id)
  401. to_node.flush_links()
  402. def handle_link_to_constant_output(self, frame_mantis_nodes, index, ui_link, to_ui_node):
  403. to_node = self.schema_nodes[(*self.tree_path_names, to_ui_node.bl_idname)]
  404. expose_when = to_node.evaluate_input('Expose when N==')
  405. # this will be None for nested Groups that have ordinary Group Input and Group Output nodes.
  406. # and in that case we just do it no matter what
  407. if expose_when is not None and expose_when > self.solve_length-1:
  408. raise GraphError(f"The Schema has a constant output at index {expose_when}"
  409. f" but it terminates at index {self.solve_length-1}. "
  410. f"The Schema must have a length of at least {expose_when+1}.")
  411. from_name = get_link_in_out(ui_link)[0]
  412. from bpy.types import NodeSocket
  413. #use it directly if it is a mantis node; this happens when the previous node was a Schema
  414. if isinstance( ui_link.from_socket, NodeSocket): # normal
  415. from_node = frame_mantis_nodes[(*self.autogen_path_names, from_name+self.index_str()) ]
  416. else: # it's a mantis socket, set manually while looping though prepped links
  417. from_node = ui_link.from_socket.node
  418. from_socket_name = ui_link.from_socket.name
  419. if from_node.node_type == 'DUMMY_SCHEMA':
  420. if isinstance( ui_link.from_socket, NodeSocket): # normal
  421. from_socket_name = ui_link.from_socket.identifier
  422. else:
  423. from_socket_name = ui_link.from_socket.name
  424. # HACK here to force it to work with ordinary node groups, which don't seem to set this value correctly.
  425. if to_ui_node.bl_idname == "NodeGroupOutput":
  426. expose_when = index # the expose-when value is None because the GroupOutput doesn't have the socket.
  427. # end HACK
  428. if index == expose_when:
  429. for outgoing in self.constant_out[ui_link.to_socket.name]:
  430. to_node = outgoing.to_node
  431. connection = from_node.outputs[from_socket_name].connect(node=to_node, socket=outgoing.to_socket)
  432. def handle_link_from_subschema_to_output(self, frame_mantis_nodes, ui_link, to_ui_node):
  433. # wire the schema node itself to the output so it will push the connections when solving
  434. to_node = self.schema_nodes[(*self.tree_path_names, to_ui_node.bl_idname)]
  435. from_name = get_link_in_out(ui_link)[0]
  436. from_node = frame_mantis_nodes[ (*self.autogen_path_names, from_name+self.index_str()) ]
  437. connection = from_node.outputs[ui_link.from_socket.identifier].connect(node=to_node, socket=ui_link.to_socket.name)
  438. def handle_link_from_array_type_to_array_out(self, original_ui_link, dummy_link):
  439. # this is so annoyingly specific lol
  440. from_node = dummy_link.nc_from; from_socket_name=dummy_link.from_socket.name
  441. for outgoing in self.array_output_connections[original_ui_link.to_socket.identifier]:
  442. to_node = outgoing.to_node
  443. connection = from_node.outputs[from_socket_name].connect(node=to_node, socket=outgoing.to_socket)
  444. # WTF is even happening here?? TODO BUG HACK
  445. 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!
  446. to_node = self.schema_nodes[(*self.tree_path_names, to_ui_node.bl_idname)] # get it by [], we want a KeyError if this fails
  447. from_name = get_link_in_out(ui_link)[0]
  448. from bpy.types import NodeSocket
  449. #use it directly if it is a mantis node; this happens when the previous node was a Schema
  450. if isinstance( ui_link.from_socket, NodeSocket): # normal
  451. from_node = frame_mantis_nodes[(*self.autogen_path_names, from_name+self.index_str()) ]
  452. else: # it's a mantis socket, set manually while looping though prepped links
  453. from_node = ui_link.from_socket.node
  454. from_socket_name = ui_link.from_socket.name
  455. for outgoing in self.array_output_connections[ui_link.to_socket.identifier]:
  456. if not from_node:
  457. from_node = self.schema_nodes[(*self.tree_path_names, from_ui_node.bl_idname)]
  458. to_node = outgoing.to_node
  459. from .schema_nodes import SchemaIndex
  460. if isinstance(from_node, SchemaIndex):
  461. signature = ("MANTIS_AUTOGENERATED", *self.tree_path_names[1:-1], self.index_str(),
  462. ui_link.from_socket.name, ui_link.from_socket.identifier)
  463. from_node = self.all_nodes.get(signature)
  464. if not from_node:
  465. from .readtree import autogen_node
  466. from_node = autogen_node(self.node.base_tree, ui_link.from_socket,
  467. signature, self.node.mContext)
  468. if ui_link.from_socket.name in ['Index']:
  469. from_node.parameters = {ui_link.from_socket.name:index}
  470. else:
  471. from_node.parameters = {ui_link.from_socket.name:self.solve_length}
  472. frame_mantis_nodes[signature]=from_node; self.solved_nodes[signature]=from_node
  473. self.all_nodes[signature]=from_node
  474. elif from_node.node_type == 'DUMMY_SCHEMA':
  475. if isinstance( ui_link.from_socket, NodeSocket): # normal
  476. from_socket_name = ui_link.from_socket.identifier
  477. else:
  478. from_socket_name = ui_link.from_socket.name
  479. if to_node.node_type == 'DUMMY_SCHEMA' and to_node.prepared:
  480. other_stem = ('SCHEMA_AUTOGENERATED', *to_node.signature[1:])
  481. other_schema_np = to_node.prototype
  482. other_schema_tree = other_schema_np.node_tree
  483. for n in other_schema_tree.nodes:
  484. if n.bl_idname not in ["SchemaArrayInput", "SchemaArrayInputGet", "SchemaArrayInputAll"]:
  485. continue
  486. out = n.outputs[outgoing.to_socket]
  487. for l in out.links:
  488. other_index_str = lambda : '.'+str(to_node.uuid)+'.'+str(index).zfill(4)
  489. other_sig = (*other_stem, l.to_node.name+other_index_str())
  490. out_node = self.all_nodes.get(other_sig)
  491. if not out_node: # if it is a sub-schema, the other node is in the parent
  492. out_node = to_node.solver.all_nodes.get(other_sig)
  493. if not out_node:
  494. raise KeyError(f"Cannot find node: {other_sig}")
  495. connection = from_node.outputs[from_socket_name].connect(node=out_node, socket=l.to_socket.name)
  496. else:
  497. connection = from_node.outputs[from_socket_name].connect(node=to_node, socket=outgoing.to_socket)
  498. def spoof_link_for_subschema_to_output_edge_case(self, ui_link, ):
  499. to_ui_node = ui_link.to_socket.node
  500. # ui_link is the link between the Schema & the output but we have mantis_links
  501. # connected up correctly. between the output and the schema's generated nodes.
  502. # so seek BACK from the output node and grab the from-node that is connected to
  503. # the link. then modify the ui_link to point to that node.
  504. from .base_definitions import DummyLink
  505. if not isinstance(ui_link, DummyLink): # make it a Dummy so i can modify it
  506. ui_link = DummyLink(ui_link.from_socket, ui_link.to_socket,
  507. multi_input_sort_id=ui_link.multi_input_sort_id)
  508. to_node = self.schema_nodes[(*self.tree_path_names, to_ui_node.bl_idname)]
  509. for inp in to_node.inputs:
  510. if inp.name == ui_link.to_socket.name: # the most recent one is from
  511. l = inp.links[-1]; break # <---- this index of the outer schema
  512. ui_link.from_nc = l.from_node; ui_link.from_socket = l.from_node.outputs[l.from_socket]
  513. return ui_link
  514. def test_is_sub_schema(self, other):
  515. for i in range(len(other.ui_signature)-1): # -1, we don't want to check this node, obviously
  516. if self.node.ui_signature[:i+1]:
  517. return False
  518. return True
  519. def prepare_nodes(self, unprepared):
  520. # At this point, we've already run a pretty exhaustive preperation phase to prep the schema's dependencies
  521. # So we should not need to add any new dependencies unless there is a bug elsewhere.
  522. # and in fact, I could skip this in some cases, and should investigate if profiling reveals a slowdown here.
  523. forbidden=set()
  524. e = None
  525. # forbid some nodes - they aren't necessary to solve the schema & cause problems.
  526. while unprepared:
  527. nc = unprepared.pop()
  528. if nc.node_type == 'DUMMY_SCHEMA' and not self.test_is_sub_schema(nc):
  529. forbidden.add(nc) # do NOT add this as a dependency.
  530. if nc in forbidden: continue # trying to resolve dependencies for.
  531. if sum([dep.prepared for dep in nc.hierarchy_dependencies]) == len(nc.hierarchy_dependencies):
  532. try:
  533. nc.bPrepare()
  534. if nc.node_type == 'DUMMY_SCHEMA':
  535. self.solve_nested_schema(nc)
  536. except Exception as e:
  537. raise execution_error_cleanup(nc, e, show_error = False)
  538. break
  539. if not nc.prepared:
  540. raise RuntimeError( f"{nc} has failed to prepare."
  541. " Please report this as a bug in mantis." )
  542. else: # Keeping this for-loop as a fallback, it should never add dependencies though
  543. can_add_me = True
  544. for dep in nc.hierarchy_dependencies:
  545. if not dep.prepared and dep not in unprepared:
  546. if dep in forbidden:
  547. can_add_me=False
  548. forbidden.add(nc) # forbid the parent, too
  549. continue
  550. unprepared.appendleft(dep)
  551. if can_add_me:
  552. unprepared.appendleft(nc) # just rotate them until they are ready.
  553. def solve_iteration(self):
  554. """ Solve an iteration of the schema.
  555. - 1 Create the Mantis Node instances that represent this iteration of the schema
  556. - 2 Connect the links from the entrypoint or previous iteration.
  557. - 3 Connect the constant and array links, and any link between nodes entirely within the tree
  558. - 4 Prepare the nodes that modify data (in case of e.g. array get index or nested schema length input)
  559. - 5 Connect the final prepared nodes
  560. and return the nodes that were created in this schema iteration (frame).
  561. This function also adds to held_links to pass data between iterations.
  562. """
  563. from .schema_nodes_ui import (SchemaIndex,
  564. SchemaArrayInput,
  565. SchemaArrayInputGet,
  566. SchemaArrayInputAll,
  567. SchemaArrayOutput,
  568. SchemaConstInput,
  569. SchemaConstOutput,
  570. SchemaOutgoingConnection,
  571. SchemaIncomingConnection,)
  572. from .utilities import clear_reroutes, link_node_containers
  573. from .base_definitions import array_output_types
  574. self.set_index_strings()
  575. frame_mantis_nodes = {}
  576. # Later we have to run bPrepare() on these guys, so we make the deque and fill it now.
  577. from collections import deque
  578. unprepared= deque()
  579. self.gen_solve_iteration_mantis_nodes(frame_mantis_nodes, unprepared)
  580. # This is where we handle node connections BETWEEN frames
  581. while(self.held_links):
  582. ui_link = self.held_links.pop()
  583. to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node
  584. if isinstance(to_ui_node, SchemaOutgoingConnection):
  585. self.handle_link_to_outgoing_connection_output(frame_mantis_nodes, ui_link)
  586. # Get the rerouted links from the graph. We don't really need to do this every iteration.
  587. # TODO: use profiling to determine if this is slow; if so: copy & reuse the data, refactor the pop()'s out.
  588. ui_links = clear_reroutes(list(self.tree.links))
  589. # Now we handle ui_links in the current frame, including those ui_links between Schema nodes and "real" nodes
  590. links_to_output = []
  591. array_input_get_link = []
  592. for ui_link in ui_links:
  593. to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node
  594. if isinstance(from_ui_node, SchemaIndex):
  595. if ui_link.from_socket.name == "Index":
  596. self.handle_link_from_index_input(self.index, frame_mantis_nodes, ui_link)
  597. elif ui_link.from_socket.name == "Schema Length":
  598. self.handle_link_from_schema_length_input(frame_mantis_nodes, ui_link)
  599. continue
  600. if isinstance(from_ui_node, SchemaIncomingConnection):
  601. if ui_link.from_socket.name in self.incoming_connections.keys():
  602. self.handle_link_from_incoming_connection_input(frame_mantis_nodes, ui_link)
  603. continue
  604. if isinstance(from_ui_node, (SchemaConstInput, NodeGroupInput)):
  605. if ui_link.from_socket.name in self.constant_in.keys():
  606. self.handle_link_from_constant_input( frame_mantis_nodes, ui_link, to_ui_node)
  607. continue
  608. if isinstance(to_ui_node, SchemaArrayInputGet):
  609. self.handle_link_to_array_input_get( frame_mantis_nodes, ui_link)
  610. continue
  611. if isinstance(from_ui_node, SchemaArrayInput):
  612. self.handle_link_from_array_input(frame_mantis_nodes, ui_link, self.index)
  613. continue
  614. if isinstance(from_ui_node, SchemaArrayInputAll):
  615. self.handle_link_from_array_input_all(frame_mantis_nodes, ui_link)
  616. continue
  617. # HOLD these links to the next iteration:
  618. if isinstance(to_ui_node, SchemaOutgoingConnection):
  619. if isinstance(from_ui_node, (MantisNodeGroup, SchemaGroup)):
  620. self.handle_link_from_subschema_to_output(frame_mantis_nodes, ui_link, to_ui_node)
  621. self.held_links.append(ui_link) # is this wise? Why am I doing this?
  622. continue
  623. # HOLD these links until prep is done a little later
  624. if isinstance(to_ui_node, (SchemaConstOutput, NodeGroupOutput)) or isinstance(to_ui_node, SchemaArrayOutput):
  625. if isinstance(from_ui_node, (MantisNodeGroup, SchemaGroup)):
  626. self.handle_link_from_subschema_to_output(frame_mantis_nodes, ui_link, to_ui_node)
  627. # both links are desirable to create, so don't continue here
  628. if from_ui_node.bl_idname in array_output_types:
  629. self.handle_link_from_subschema_to_output(frame_mantis_nodes, ui_link, to_ui_node)
  630. # this one wires links around - we need to finish connecting it to the output
  631. # before we can prepare it. Otherwise, it will send a link from *itself* instead of
  632. # rewiring one of its *inputs* to the next node.
  633. # We have to prep it, and then deal with the links between it and the array. pain.
  634. links_to_output.append(ui_link)
  635. continue
  636. if isinstance(from_ui_node, SchemaArrayInputGet):
  637. array_input_get_link.append(ui_link)
  638. continue
  639. # for any of the special cases, we hit a 'continue' block. So this connection is not special, and is made here.
  640. connection = link_node_containers(self.autogen_path_names, ui_link,
  641. frame_mantis_nodes, from_suffix=self.index_str(),
  642. to_suffix=self.index_str())
  643. for signature, node in frame_mantis_nodes.items():
  644. self.solved_nodes[signature]=node
  645. if node.node_type == "DUMMY_SCHEMA":
  646. # make sure to add the nodes to the group's sockets if the user set them directly
  647. from .readtree import make_connections_to_ng_dummy
  648. make_connections_to_ng_dummy(
  649. self.node.base_tree,
  650. self.autogen_path_names,
  651. {}, # just pass an empty dict, this argument is not needed in this context
  652. self.all_nodes,
  653. node)
  654. from .utilities import init_schema_dependencies
  655. init_schema_dependencies(node, self.all_nodes)
  656. else:
  657. init_dependencies(node) # it is hard to overstate how important this single line of code is
  658. # We have to prepare the nodes leading to Schema Length
  659. unprepared=deque()
  660. for node in frame_mantis_nodes.values():
  661. if node.node_type == 'DUMMY_SCHEMA' and (schema_len_in := node.inputs.get("Schema Length")):
  662. for l in schema_len_in.links:
  663. unprepared.append(l.from_node)
  664. self.prepare_nodes(unprepared)
  665. # We have to prepare the nodes leading to Array Input Get
  666. for ui_link in array_input_get_link:
  667. from_name = get_link_in_out(ui_link)[0]
  668. # because this both provides and receives deps, it must be solved first.
  669. from_node = self.schema_nodes.get( (*self.node.ui_signature, ui_link.from_node.bl_idname) )
  670. self.handle_link_from_array_input_get(frame_mantis_nodes, ui_link )
  671. # Finally, we have to prepare nodes leading to outputs.
  672. for i in range(len(links_to_output)):
  673. unprepared=deque()
  674. ui_link = links_to_output[i]
  675. to_ui_node = ui_link.to_socket.node; from_ui_node = ui_link.from_socket.node
  676. # ugly workaround here in a very painful edge case...
  677. if isinstance(from_ui_node, (MantisNodeGroup, SchemaGroup)) or\
  678. from_ui_node.bl_idname in array_output_types:
  679. ui_link=self.spoof_link_for_subschema_to_output_edge_case(ui_link)
  680. links_to_output[i] = ui_link
  681. from_name = get_link_in_out(ui_link)[0]
  682. signature = (*self.autogen_path_names, from_name+self.index_str())
  683. #use it directly if it is a mantis node; this happens when the previous node was a Schema
  684. if hasattr(ui_link, "from_node") and (from_node := self.schema_nodes.get( (*self.node.ui_signature, ui_link.from_node.bl_idname))):
  685. unprepared.append(from_node)
  686. unprepared.extend(from_node.hierarchy_dependencies)
  687. elif from_node := frame_mantis_nodes.get(signature):
  688. unprepared.append(from_node)
  689. unprepared.extend(from_node.hierarchy_dependencies)
  690. else:
  691. raise RuntimeError(" 671 there has been an error parsing the tree. Please report this as a bug.")
  692. self.prepare_nodes(unprepared) # prepare only the dependencies we need for this link
  693. # and handle the output by the specific type
  694. if isinstance(to_ui_node, (SchemaConstOutput, NodeGroupOutput)):
  695. self.handle_link_to_constant_output(frame_mantis_nodes, self.index, ui_link, to_ui_node)
  696. if isinstance(to_ui_node, SchemaArrayOutput):
  697. if from_node.bl_idname in array_output_types:
  698. from .base_definitions import DummyLink
  699. for l in from_node.rerouted:
  700. new_link = DummyLink(
  701. from_socket=l.from_node.outputs[l.from_socket],
  702. to_socket=l.to_node.inputs[l.to_socket],
  703. nc_from=l.from_node, nc_to=l.to_node,
  704. multi_input_sort_id=l.multi_input_sort_id
  705. )
  706. self.handle_link_from_array_type_to_array_out(ui_link, new_link)
  707. else:
  708. self.handle_link_to_array_output(frame_mantis_nodes, self.index, ui_link, to_ui_node, from_ui_node)
  709. return frame_mantis_nodes
  710. def solve_nested_schema(self, schema_nc):
  711. """ Solves all schema node groups found in this Schema. This is a recursive function, which will
  712. solve all levels of nested schema - since this function is called by solver.solve().
  713. """
  714. solver=None
  715. if schema_nc.prepared == False:
  716. all_nodes = self.all_nodes.copy()
  717. ui_node = schema_nc.prototype
  718. length = schema_nc.evaluate_input("Schema Length")
  719. tree = ui_node.node_tree
  720. if schema_nc.prototype.bl_idname == "MantisNodeGroup":
  721. prOrange(f"Expanding Node Group {tree.name} in node {schema_nc}.")
  722. else:
  723. prOrange(f"Expanding schema {tree.name} in node {schema_nc} with length {length}.")
  724. solver = SchemaSolver(schema_nc, all_nodes, ui_node, schema_nc.ui_signature, error_popups=self.error_popups)
  725. solved_nodes = solver.solve()
  726. schema_nc.prepared = True
  727. for k,v in solved_nodes.items():
  728. self.solved_nodes[k]=v
  729. return solver
  730. def finalize(self, frame_nc):
  731. from .schema_nodes_ui import (SchemaOutgoingConnection,)
  732. for i in range(len(self.held_links)):
  733. link = self.held_links.pop()
  734. to_np = link.to_socket.node; from_np = link.from_socket.node
  735. if isinstance(to_np, SchemaOutgoingConnection):
  736. if link.to_socket.name in self.outgoing_connections.keys():
  737. if (outgoing_links := self.outgoing_connections[link.to_socket.name]) is None: continue
  738. for outgoing in outgoing_links:
  739. if outgoing:
  740. to_node = outgoing.to_node
  741. from_node =frame_nc[(*self.autogen_path_names, from_np.name+self.index_str()) ]
  742. from_socket_name = link.from_socket.name
  743. if from_node.node_type in ['DUMMY_SCHEMA']:
  744. from_socket_name = link.from_socket.identifier
  745. connection = from_node.outputs[from_socket_name].connect(node=to_node, socket=outgoing.to_socket)
  746. # we need to kill the link between the Schema itself and the next node and update the deps. Otherwise:confusing bugs.
  747. outgoing.die(); init_dependencies(to_node)
  748. # else: # the node just isn't connected out this socket.
  749. # # solve all unsolved nested schemas...
  750. for schema_sig, schema_nc in self.nested_schemas.items():
  751. self.solve_nested_schema(schema_nc)
  752. for n in self.autogenerated_nodes.values():
  753. init_connections(n)
  754. for c in n.connections:
  755. init_dependencies(c)
  756. all_outgoing_links = []
  757. for conn in self.outgoing_connections.values():
  758. for outgoing in conn:
  759. all_outgoing_links.append(outgoing)
  760. for conn in self.constant_out.values():
  761. for outgoing in conn:
  762. all_outgoing_links.append(outgoing)
  763. for conn in self.array_output_connections.values():
  764. for outgoing in conn:
  765. all_outgoing_links.append(outgoing)
  766. for outgoing in all_outgoing_links:
  767. to_node = outgoing.to_node
  768. for l in to_node.inputs[outgoing.to_socket].links:
  769. if self.node == l.from_node:
  770. l.die()
  771. for inp in self.node.inputs.values():
  772. for l in inp.links:
  773. init_connections(l.from_node) # to force it to have hierarchy connections with the new nodes.
  774. def solve(self):
  775. if self.solve_length < 1:
  776. from .base_definitions import GraphError
  777. for o in self.node.outputs:
  778. if o.is_linked:
  779. raise GraphError(f"ERROR: Schema {self} has a length"
  780. " of 0 but other nodes depend on it.")
  781. print (f"WARN: Schema {self} has a length of 0 or less and will not expand.")
  782. return {} # just don't do anything - it's OK to have a noop schema if it doesn't have dependencies.
  783. for index in range(self.solve_length):
  784. self.index = index
  785. frame_mantis_nodes = self.solve_iteration()
  786. for sig, nc in frame_mantis_nodes.items():
  787. if nc.node_type == 'DUMMY_SCHEMA':
  788. self.nested_schemas[sig] = nc
  789. self.finalize(frame_mantis_nodes)
  790. self.solved = True
  791. self.node.prepared = True
  792. prGreen(f"Schema declared {len(self.solved_nodes)} nodes.\n")
  793. return self.solved_nodes