compiled.py 3.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. from modelverse_kernel.primitives import PrimitiveFinished
  2. import modelverse_jit.runtime as jit_runtime
  3. def reverseKeyLookupMulti(a, b, **remainder):
  4. edges, b_val, result = yield [("RO", [a]), ("RV", [b]), ("CN", [])]
  5. expanded_edges = yield [("RE", [i]) for i in edges]
  6. values = yield [("RV", [i[1]]) for i in expanded_edges]
  7. # Keep results in a local Python set, as we want to bundle as many requests as possible
  8. todo = set()
  9. for i, edge in enumerate(values):
  10. if b_val == edge:
  11. todo.add(i)
  12. outgoings = yield [("RO", [edges[i]]) for i in todo]
  13. values = yield [("RE", [outgoing[0]]) for outgoing in outgoings]
  14. yield [("CE", [result, value[1]]) for value in values]
  15. raise PrimitiveFinished(result)
  16. def reverseKeyLookup(a, b, **remainder):
  17. edges_out, edges_in = yield [("RO", [a]), ("RI", [b])]
  18. options = set(edges_out) & set(edges_in)
  19. if options:
  20. # Select one option randomly
  21. edge = options.pop()
  22. out_edges, = yield [("RO", [edge])]
  23. # Select one option randomly
  24. out_edge = out_edges.pop()
  25. e, = yield [("RE", [out_edge])]
  26. s, t = e
  27. raise PrimitiveFinished(t)
  28. else:
  29. result, = yield [("CNV", ["(unknown: %s)" % b])]
  30. raise PrimitiveFinished(result)
  31. def set_copy(a, **remainder):
  32. b, = yield [("CN", [])]
  33. links, = yield [("RO", [a])]
  34. exp_links = yield [("RE", [i]) for i in links]
  35. _ = yield [("CE", [b, i[1]]) for i in exp_links]
  36. raise PrimitiveFinished(b)
  37. def instantiated_name(a, b, **remainder):
  38. name_value, = yield [("RV", [b])]
  39. if name_value == "":
  40. b, = yield [("CNV", ["__" + str(a)])]
  41. raise PrimitiveFinished(b)
  42. def set_merge(a, b, **remainder):
  43. outputs, = yield [("RO", [b])]
  44. values = yield [("RE", [i]) for i in outputs]
  45. yield [("CE", [a, i[1]]) for i in values]
  46. raise PrimitiveFinished(a)
  47. def has_value(a, **remainder):
  48. v, = yield [("RV", [a])]
  49. if v is None:
  50. result, = yield [("CNV", [False])]
  51. else:
  52. result, = yield [("CNV", [True])]
  53. raise PrimitiveFinished(result)
  54. def make_reverse_dictionary(a, **remainder):
  55. reverse, = yield [("CN", [])]
  56. key_nodes, = yield [("RDK", [a])]
  57. values = yield [("RDN", [a, i]) for i in key_nodes]
  58. yield [("CD", [reverse, str(v), k]) for k, v in zip(key_nodes, values)]
  59. raise PrimitiveFinished(reverse)
  60. def dict_eq(a, b, **remainder):
  61. key_nodes, = yield [("RDK", [a])]
  62. key_values = yield [("RV", [i]) for i in key_nodes]
  63. values = yield [("RD", [a, i]) for i in key_values]
  64. values = yield [("RV", [i]) for i in values]
  65. a_dict = dict(zip(key_values, values))
  66. key_nodes, = yield [("RDK", [b])]
  67. key_values = yield [("RV", [i]) for i in key_nodes]
  68. values = yield [("RD", [b, i]) for i in key_values]
  69. values = yield [("RV", [i]) for i in values]
  70. b_dict = dict(zip(key_values, values))
  71. result, = yield [("CNV", [a_dict == b_dict])]
  72. raise PrimitiveFinished(result)
  73. def string_substr(a, b, c, **remainder):
  74. a_val, b_val, c_val = yield [("RV", [a]),
  75. ("RV", [b]),
  76. ("RV", [c])]
  77. try:
  78. new_value = a_val[b_val:c_val]
  79. except:
  80. new_value = ""
  81. result, = yield [("CNV", [new_value])]
  82. raise PrimitiveFinished(result)