generateQuadTree.js 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461
  1. const createPatternBuilder = require('./createPatternBuilder');
  2. const getVariableName = require('./getVariableName');
  3. module.exports = generateQuadTreeFunction;
  4. module.exports.generateQuadTreeFunctionBody = generateQuadTreeFunctionBody;
  5. // These exports are for InlineTransform tool.
  6. // InlineTransform: getInsertStackCode
  7. module.exports.getInsertStackCode = getInsertStackCode;
  8. // InlineTransform: getQuadNodeCode
  9. module.exports.getQuadNodeCode = getQuadNodeCode;
  10. // InlineTransform: isSamePosition
  11. module.exports.isSamePosition = isSamePosition;
  12. // InlineTransform: getChildBodyCode
  13. module.exports.getChildBodyCode = getChildBodyCode;
  14. // InlineTransform: setChildBodyCode
  15. module.exports.setChildBodyCode = setChildBodyCode;
  16. function generateQuadTreeFunction(dimension) {
  17. let code = generateQuadTreeFunctionBody(dimension);
  18. return (new Function(code))();
  19. }
  20. function generateQuadTreeFunctionBody(dimension) {
  21. let pattern = createPatternBuilder(dimension);
  22. let quadCount = Math.pow(2, dimension);
  23. let code = `
  24. ${getInsertStackCode()}
  25. ${getQuadNodeCode(dimension)}
  26. ${isSamePosition(dimension)}
  27. ${getChildBodyCode(dimension)}
  28. ${setChildBodyCode(dimension)}
  29. function createQuadTree(options, random) {
  30. options = options || {};
  31. options.gravity = typeof options.gravity === 'number' ? options.gravity : -1;
  32. options.theta = typeof options.theta === 'number' ? options.theta : 0.8;
  33. var gravity = options.gravity;
  34. var updateQueue = [];
  35. var insertStack = new InsertStack();
  36. var theta = options.theta;
  37. var nodesCache = [];
  38. var currentInCache = 0;
  39. var root = newNode();
  40. return {
  41. insertBodies: insertBodies,
  42. /**
  43. * Gets root node if it is present
  44. */
  45. getRoot: function() {
  46. return root;
  47. },
  48. updateBodyForce: update,
  49. options: function(newOptions) {
  50. if (newOptions) {
  51. if (typeof newOptions.gravity === 'number') {
  52. gravity = newOptions.gravity;
  53. }
  54. if (typeof newOptions.theta === 'number') {
  55. theta = newOptions.theta;
  56. }
  57. return this;
  58. }
  59. return {
  60. gravity: gravity,
  61. theta: theta
  62. };
  63. }
  64. };
  65. function newNode() {
  66. // To avoid pressure on GC we reuse nodes.
  67. var node = nodesCache[currentInCache];
  68. if (node) {
  69. ${assignQuads(' node.')}
  70. node.body = null;
  71. node.mass = ${pattern('node.mass_{var} = ', {join: ''})}0;
  72. ${pattern('node.min_{var} = node.max_{var} = ', {join: ''})}0;
  73. } else {
  74. node = new QuadNode();
  75. nodesCache[currentInCache] = node;
  76. }
  77. ++currentInCache;
  78. return node;
  79. }
  80. function update(sourceBody) {
  81. var queue = updateQueue;
  82. var v;
  83. ${pattern('var d{var};', {indent: 4})}
  84. var r;
  85. ${pattern('var f{var} = 0;', {indent: 4})}
  86. var queueLength = 1;
  87. var shiftIdx = 0;
  88. var pushIdx = 1;
  89. queue[0] = root;
  90. while (queueLength) {
  91. var node = queue[shiftIdx];
  92. var body = node.body;
  93. queueLength -= 1;
  94. shiftIdx += 1;
  95. var differentBody = (body !== sourceBody);
  96. if (body && differentBody) {
  97. // If the current node is a leaf node (and it is not source body),
  98. // calculate the force exerted by the current node on body, and add this
  99. // amount to body's net force.
  100. ${pattern('d{var} = body.pos.{var} - sourceBody.pos.{var};', {indent: 8})}
  101. r = Math.sqrt(${pattern('d{var} * d{var}', {join: ' + '})});
  102. if (r === 0) {
  103. // Poor man's protection against zero distance.
  104. ${pattern('d{var} = (random.nextDouble() - 0.5) / 50;', {indent: 10})}
  105. r = Math.sqrt(${pattern('d{var} * d{var}', {join: ' + '})});
  106. }
  107. // This is standard gravitation force calculation but we divide
  108. // by r^3 to save two operations when normalizing force vector.
  109. v = gravity * body.mass * sourceBody.mass / (r * r * r);
  110. ${pattern('f{var} += v * d{var};', {indent: 8})}
  111. } else if (differentBody) {
  112. // Otherwise, calculate the ratio s / r, where s is the width of the region
  113. // represented by the internal node, and r is the distance between the body
  114. // and the node's center-of-mass
  115. ${pattern('d{var} = node.mass_{var} / node.mass - sourceBody.pos.{var};', {indent: 8})}
  116. r = Math.sqrt(${pattern('d{var} * d{var}', {join: ' + '})});
  117. if (r === 0) {
  118. // Sorry about code duplication. I don't want to create many functions
  119. // right away. Just want to see performance first.
  120. ${pattern('d{var} = (random.nextDouble() - 0.5) / 50;', {indent: 10})}
  121. r = Math.sqrt(${pattern('d{var} * d{var}', {join: ' + '})});
  122. }
  123. // If s / r < θ, treat this internal node as a single body, and calculate the
  124. // force it exerts on sourceBody, and add this amount to sourceBody's net force.
  125. if ((node.max_${getVariableName(0)} - node.min_${getVariableName(0)}) / r < theta) {
  126. // in the if statement above we consider node's width only
  127. // because the region was made into square during tree creation.
  128. // Thus there is no difference between using width or height.
  129. v = gravity * node.mass * sourceBody.mass / (r * r * r);
  130. ${pattern('f{var} += v * d{var};', {indent: 10})}
  131. } else {
  132. // Otherwise, run the procedure recursively on each of the current node's children.
  133. // I intentionally unfolded this loop, to save several CPU cycles.
  134. ${runRecursiveOnChildren()}
  135. }
  136. }
  137. }
  138. ${pattern('sourceBody.force.{var} += f{var};', {indent: 4})}
  139. }
  140. function insertBodies(bodies) {
  141. ${pattern('var {var}min = Number.MAX_VALUE;', {indent: 4})}
  142. ${pattern('var {var}max = Number.MIN_VALUE;', {indent: 4})}
  143. var i = bodies.length;
  144. // To reduce quad tree depth we are looking for exact bounding box of all particles.
  145. while (i--) {
  146. var pos = bodies[i].pos;
  147. ${pattern('if (pos.{var} < {var}min) {var}min = pos.{var};', {indent: 6})}
  148. ${pattern('if (pos.{var} > {var}max) {var}max = pos.{var};', {indent: 6})}
  149. }
  150. // Makes the bounds square.
  151. var maxSideLength = -Infinity;
  152. ${pattern('if ({var}max - {var}min > maxSideLength) maxSideLength = {var}max - {var}min ;', {indent: 4})}
  153. currentInCache = 0;
  154. root = newNode();
  155. ${pattern('root.min_{var} = {var}min;', {indent: 4})}
  156. ${pattern('root.max_{var} = {var}min + maxSideLength;', {indent: 4})}
  157. i = bodies.length - 1;
  158. if (i >= 0) {
  159. root.body = bodies[i];
  160. }
  161. while (i--) {
  162. insert(bodies[i], root);
  163. }
  164. }
  165. function insert(newBody) {
  166. insertStack.reset();
  167. insertStack.push(root, newBody);
  168. while (!insertStack.isEmpty()) {
  169. var stackItem = insertStack.pop();
  170. var node = stackItem.node;
  171. var body = stackItem.body;
  172. if (!node.body) {
  173. // This is internal node. Update the total mass of the node and center-of-mass.
  174. ${pattern('var {var} = body.pos.{var};', {indent: 8})}
  175. node.mass += body.mass;
  176. ${pattern('node.mass_{var} += body.mass * {var};', {indent: 8})}
  177. // Recursively insert the body in the appropriate quadrant.
  178. // But first find the appropriate quadrant.
  179. var quadIdx = 0; // Assume we are in the 0's quad.
  180. ${pattern('var min_{var} = node.min_{var};', {indent: 8})}
  181. ${pattern('var max_{var} = (min_{var} + node.max_{var}) / 2;', {indent: 8})}
  182. ${assignInsertionQuadIndex(8)}
  183. var child = getChild(node, quadIdx);
  184. if (!child) {
  185. // The node is internal but this quadrant is not taken. Add
  186. // subnode to it.
  187. child = newNode();
  188. ${pattern('child.min_{var} = min_{var};', {indent: 10})}
  189. ${pattern('child.max_{var} = max_{var};', {indent: 10})}
  190. child.body = body;
  191. setChild(node, quadIdx, child);
  192. } else {
  193. // continue searching in this quadrant.
  194. insertStack.push(child, body);
  195. }
  196. } else {
  197. // We are trying to add to the leaf node.
  198. // We have to convert current leaf into internal node
  199. // and continue adding two nodes.
  200. var oldBody = node.body;
  201. node.body = null; // internal nodes do not cary bodies
  202. if (isSamePosition(oldBody.pos, body.pos)) {
  203. // Prevent infinite subdivision by bumping one node
  204. // anywhere in this quadrant
  205. var retriesCount = 3;
  206. do {
  207. var offset = random.nextDouble();
  208. ${pattern('var d{var} = (node.max_{var} - node.min_{var}) * offset;', {indent: 12})}
  209. ${pattern('oldBody.pos.{var} = node.min_{var} + d{var};', {indent: 12})}
  210. retriesCount -= 1;
  211. // Make sure we don't bump it out of the box. If we do, next iteration should fix it
  212. } while (retriesCount > 0 && isSamePosition(oldBody.pos, body.pos));
  213. if (retriesCount === 0 && isSamePosition(oldBody.pos, body.pos)) {
  214. // This is very bad, we ran out of precision.
  215. // if we do not return from the method we'll get into
  216. // infinite loop here. So we sacrifice correctness of layout, and keep the app running
  217. // Next layout iteration should get larger bounding box in the first step and fix this
  218. return;
  219. }
  220. }
  221. // Next iteration should subdivide node further.
  222. insertStack.push(node, oldBody);
  223. insertStack.push(node, body);
  224. }
  225. }
  226. }
  227. }
  228. return createQuadTree;
  229. `;
  230. return code;
  231. function assignInsertionQuadIndex(indentCount) {
  232. let insertionCode = [];
  233. let indent = Array(indentCount + 1).join(' ');
  234. for (let i = 0; i < dimension; ++i) {
  235. insertionCode.push(indent + `if (${getVariableName(i)} > max_${getVariableName(i)}) {`);
  236. insertionCode.push(indent + ` quadIdx = quadIdx + ${Math.pow(2, i)};`);
  237. insertionCode.push(indent + ` min_${getVariableName(i)} = max_${getVariableName(i)};`);
  238. insertionCode.push(indent + ` max_${getVariableName(i)} = node.max_${getVariableName(i)};`);
  239. insertionCode.push(indent + `}`);
  240. }
  241. return insertionCode.join('\n');
  242. // if (x > max_x) { // somewhere in the eastern part.
  243. // quadIdx = quadIdx + 1;
  244. // left = right;
  245. // right = node.right;
  246. // }
  247. }
  248. function runRecursiveOnChildren() {
  249. let indent = Array(11).join(' ');
  250. let recursiveCode = [];
  251. for (let i = 0; i < quadCount; ++i) {
  252. recursiveCode.push(indent + `if (node.quad${i}) {`);
  253. recursiveCode.push(indent + ` queue[pushIdx] = node.quad${i};`);
  254. recursiveCode.push(indent + ` queueLength += 1;`);
  255. recursiveCode.push(indent + ` pushIdx += 1;`);
  256. recursiveCode.push(indent + `}`);
  257. }
  258. return recursiveCode.join('\n');
  259. // if (node.quad0) {
  260. // queue[pushIdx] = node.quad0;
  261. // queueLength += 1;
  262. // pushIdx += 1;
  263. // }
  264. }
  265. function assignQuads(indent) {
  266. // this.quad0 = null;
  267. // this.quad1 = null;
  268. // this.quad2 = null;
  269. // this.quad3 = null;
  270. let quads = [];
  271. for (let i = 0; i < quadCount; ++i) {
  272. quads.push(`${indent}quad${i} = null;`);
  273. }
  274. return quads.join('\n');
  275. }
  276. }
  277. function isSamePosition(dimension) {
  278. let pattern = createPatternBuilder(dimension);
  279. return `
  280. function isSamePosition(point1, point2) {
  281. ${pattern('var d{var} = Math.abs(point1.{var} - point2.{var});', {indent: 2})}
  282. return ${pattern('d{var} < 1e-8', {join: ' && '})};
  283. }
  284. `;
  285. }
  286. function setChildBodyCode(dimension) {
  287. var quadCount = Math.pow(2, dimension);
  288. return `
  289. function setChild(node, idx, child) {
  290. ${setChildBody()}
  291. }`;
  292. function setChildBody() {
  293. let childBody = [];
  294. for (let i = 0; i < quadCount; ++i) {
  295. let prefix = (i === 0) ? ' ' : ' else ';
  296. childBody.push(`${prefix}if (idx === ${i}) node.quad${i} = child;`);
  297. }
  298. return childBody.join('\n');
  299. // if (idx === 0) node.quad0 = child;
  300. // else if (idx === 1) node.quad1 = child;
  301. // else if (idx === 2) node.quad2 = child;
  302. // else if (idx === 3) node.quad3 = child;
  303. }
  304. }
  305. function getChildBodyCode(dimension) {
  306. return `function getChild(node, idx) {
  307. ${getChildBody()}
  308. return null;
  309. }`;
  310. function getChildBody() {
  311. let childBody = [];
  312. let quadCount = Math.pow(2, dimension);
  313. for (let i = 0; i < quadCount; ++i) {
  314. childBody.push(` if (idx === ${i}) return node.quad${i};`);
  315. }
  316. return childBody.join('\n');
  317. // if (idx === 0) return node.quad0;
  318. // if (idx === 1) return node.quad1;
  319. // if (idx === 2) return node.quad2;
  320. // if (idx === 3) return node.quad3;
  321. }
  322. }
  323. function getQuadNodeCode(dimension) {
  324. let pattern = createPatternBuilder(dimension);
  325. let quadCount = Math.pow(2, dimension);
  326. var quadNodeCode = `
  327. function QuadNode() {
  328. // body stored inside this node. In quad tree only leaf nodes (by construction)
  329. // contain bodies:
  330. this.body = null;
  331. // Child nodes are stored in quads. Each quad is presented by number:
  332. // 0 | 1
  333. // -----
  334. // 2 | 3
  335. ${assignQuads(' this.')}
  336. // Total mass of current node
  337. this.mass = 0;
  338. // Center of mass coordinates
  339. ${pattern('this.mass_{var} = 0;', {indent: 2})}
  340. // bounding box coordinates
  341. ${pattern('this.min_{var} = 0;', {indent: 2})}
  342. ${pattern('this.max_{var} = 0;', {indent: 2})}
  343. }
  344. `;
  345. return quadNodeCode;
  346. function assignQuads(indent) {
  347. // this.quad0 = null;
  348. // this.quad1 = null;
  349. // this.quad2 = null;
  350. // this.quad3 = null;
  351. let quads = [];
  352. for (let i = 0; i < quadCount; ++i) {
  353. quads.push(`${indent}quad${i} = null;`);
  354. }
  355. return quads.join('\n');
  356. }
  357. }
  358. function getInsertStackCode() {
  359. return `
  360. /**
  361. * Our implementation of QuadTree is non-recursive to avoid GC hit
  362. * This data structure represent stack of elements
  363. * which we are trying to insert into quad tree.
  364. */
  365. function InsertStack () {
  366. this.stack = [];
  367. this.popIdx = 0;
  368. }
  369. InsertStack.prototype = {
  370. isEmpty: function() {
  371. return this.popIdx === 0;
  372. },
  373. push: function (node, body) {
  374. var item = this.stack[this.popIdx];
  375. if (!item) {
  376. // we are trying to avoid memory pressure: create new element
  377. // only when absolutely necessary
  378. this.stack[this.popIdx] = new InsertStackElement(node, body);
  379. } else {
  380. item.node = node;
  381. item.body = body;
  382. }
  383. ++this.popIdx;
  384. },
  385. pop: function () {
  386. if (this.popIdx > 0) {
  387. return this.stack[--this.popIdx];
  388. }
  389. },
  390. reset: function () {
  391. this.popIdx = 0;
  392. }
  393. };
  394. function InsertStackElement(node, body) {
  395. this.node = node; // QuadTree node
  396. this.body = body; // physical body which needs to be inserted to node
  397. }
  398. `;
  399. }