Fix getDefinition not working for modified/new procedures

core/procedures.js

@@ -394,8 +394,10 @@ Blockly.Procedures.mutateCallers = function(defBlock) {
  * @return {Blockly.Block} The procedure definition block, or null not found.
  */
 Blockly.Procedures.getDefinition = function(name, workspace) {
-  // Assume that a procedure definition is a top block.
-  var blocks = workspace.getTopBlocks(false);
+  // Do not assume procedure is a top block. Some languages allow nested
+  // procedures. Also do not assume it is one of the built-in blocks. Only
+  // rely on getProcedureDef.
+  var blocks = workspace.getAllBlocks(false);
   for (var i = 0; i < blocks.length; i++) {
     if (blocks[i].getProcedureDef) {
       var procedureBlock = /** @type {!Blockly.Procedures.ProcedureBlock} */ (

tests/mocha/procedures_test.js

@@ -18,7 +18,7 @@ suite('Procedures', function() {
   teardown(function() {
     sharedTestTeardown.call(this);
   });
-  
+
   suite('allProcedures', function() {
     test('Only Procedures', function() {
       var noReturnBlock = new Blockly.Block(this.workspace, 'procedures_defnoreturn');
@@ -367,6 +367,51 @@ suite('Procedures', function() {
     });
   });
 
+  suite('getDefinition - Modified cases', function() {
+    setup(function() {
+      Blockly.Blocks['new_proc'] = {
+        init: function() { },
+        getProcedureDef: function() {
+          return [this.name, [], false];
+        },
+        name: 'test'
+      };
+
+      Blockly.Blocks['nested_proc'] = {
+        init: function() {
+          this.setPreviousStatement(true, null);
+          this.setNextStatement(true, null);
+        },
+        getProcedureDef: function() {
+          return [this.name, [], false];
+        },
+        name: 'test',
+      };
+    });
+
+    teardown(function() {
+      delete Blockly.Blocks['new_proc'];
+      delete Blockly.Blocks['nested_proc'];
+    });
+
+    test('Custom procedure block', function() {
+      // Do not require procedures to be the built-in procedures.
+      var defBlock = new Blockly.Block(this.workspace, 'new_proc');
+      var def = Blockly.Procedures.getDefinition('test', this.workspace);
+      chai.assert.equal(def, defBlock);
+    });
+
+    test('Stacked procedures', function() {
+      var blockA = new Blockly.Block(this.workspace, 'nested_proc');
+      var blockB = new Blockly.Block(this.workspace, 'nested_proc');
+      blockA.name = 'a';
+      blockB.name = 'b';
+      blockA.nextConnection.connect(blockB.previousConnection);
+      var def = Blockly.Procedures.getDefinition('b', this.workspace);
+      chai.assert.equal(def, blockB);
+    });
+  });
+
   const testSuites = [
     {title: 'procedures_defreturn', hasReturn: true,
       defType: 'procedures_defreturn', callType: 'procedures_callreturn'},