How to access the return statement of a function in Javascript

Question

I'm looking for a way to get the return expression of a function in Javascript. I'm expecting to get it as a string but any type would do. Also, I'm looking for a way to set the return statement of a function without affecting the rest of it.

Example:

let f1 = () => {
    let x = 1
    return 2 + x
}

let f2 = () => {
    let y = 2
    return 3 + y
}

let s = f1.getReturnExpression()
console.log(s) // 2 + x

f2.setReturnExpression(s)
console.log(f2.toString()) // "function f2() {  let y = 2; return 2 + x; }

Answer 1

Extracting the return statements should be "easy" but I am not sure how you can reconstruct a new function and assign it to an existing function at runtime without damaging scope somehow.

Let's say that extracting return statements is enough for now. I would definitely not recommend using regular expressions or strings manipulation for that; you'll just get it wrong.

You should use tools like esprima that allow you to deconstruct a piece of code into a data structure that you can inspect.

Given the following function:

function foo(x) {
  if (x) {
    return x + 10;
  } else {
    return x - 10;
  }
}

Pass it to esprima:

esprima.tokenize(foo.toString());

This is what you get:

[ { type: 'Keyword', value: 'function' },
  { type: 'Identifier', value: 'foo' },
  { type: 'Punctuator', value: '(' },
  { type: 'Identifier', value: 'x' },
  { type: 'Punctuator', value: ')' },
  { type: 'Punctuator', value: '{' },
  { type: 'Keyword', value: 'if' },
  { type: 'Punctuator', value: '(' },
  { type: 'Identifier', value: 'x' },
  { type: 'Punctuator', value: ')' },
  { type: 'Punctuator', value: '{' },
  { type: 'Keyword', value: 'return' },
  { type: 'Identifier', value: 'x' },
  { type: 'Punctuator', value: '+' },
  { type: 'Numeric', value: '10' },
  { type: 'Punctuator', value: ';' },
  { type: 'Punctuator', value: '}' },
  { type: 'Keyword', value: 'else' },
  { type: 'Punctuator', value: '{' },
  { type: 'Keyword', value: 'return' },
  { type: 'Identifier', value: 'x' },
  { type: 'Punctuator', value: '-' },
  { type: 'Numeric', value: '10' },
  { type: 'Punctuator', value: ';' },
  { type: 'Punctuator', value: '}' },
  { type: 'Punctuator', value: '}' } ]

This is function foo in a JavaScript data structure that you can manipulate in ... JavaScript...
A m a z i n g 🤓

Now it should be fairly obvious to work out how to extract the two possible return statements.
However, it is important to note that I do assume that each return statement ends with a semicolon ;.

To identify the start of a return statement we can write a isReturnStatement function:

const isReturnStatement = token =>
  token.type === 'Keyword' &&
  token.value === 'return';

And to identify the end of such statement, we can write a isEndOfStatement function:

const isEndOfStatement = token =>
  token.type === 'Punctuator' &&
  token.value === ';';

As you can see you're now dealing with a JavaScript data structure, an array of tokens. Each token representing a piece of code syntax.

You now need to identify the start of a return statement and its end. Save that and continue the process until you have inspected the entire array.

Here's one way to do it:

function foo(x) {
  if (x) {
    return x + 10;
  } else {
    return x - 10;
  }
}

const ast = esprima.tokenize(foo.toString());

console.log(

  getReturnStatements(ast)
  
);

<script src="https://cdnjs.cloudflare.com/ajax/libs/esprima/2.7.3/esprima.min.js"></script>

<script>

const isReturnStatement = token =>
  token.type === 'Keyword' &&
  token.value === 'return';

const isEndOfStatement = token =>
  token.type === 'Punctuator' &&
  token.value === ';';

const getReturnStatements = ast => {
  const scan = ast.reduce((search, token) => {
    if (isReturnStatement(token)) {
      search.capture = true;
      search.stmt.push('');
    } else if (search.capture && !isEndOfStatement(token)) {
      search.stmt[search.stmt.length - 1] += token.value;
    } else {
      search.capture = false;
    }
    return search;
  }, {capture: false, stmt: []});
  return scan.stmt;
};

</script>

Note that for improved accuracy you should be parsing code instead of tokenizing it but I hope this gives you a good enough start.