wrap.js 1.43 KB
import types from './types/index.js';
import BlockStatement from './BlockStatement.js';
import Node from './Node.js';
import keys from './keys.js';

const statementsWithBlocks = {
	IfStatement: 'consequent',
	ForStatement: 'body',
	ForInStatement: 'body',
	ForOfStatement: 'body',
	WhileStatement: 'body',
	DoWhileStatement: 'body',
	ArrowFunctionExpression: 'body'
};

export default function wrap ( raw, parent ) {
	if ( !raw ) return;

	if ( 'length' in raw ) {
		let i = raw.length;
		while ( i-- ) wrap( raw[i], parent );
		return;
	}

	// with e.g. shorthand properties, key and value are
	// the same node. We don't want to wrap an object twice
	if ( raw.__wrapped ) return;
	raw.__wrapped = true;

	if ( !keys[ raw.type ] ) {
		keys[ raw.type ] = Object.keys( raw ).filter( key => typeof raw[ key ] === 'object' );
	}

	// special case – body-less if/for/while statements. TODO others?
	const bodyType = statementsWithBlocks[ raw.type ];
	if ( bodyType && raw[ bodyType ].type !== 'BlockStatement' ) {
		const expression = raw[ bodyType ];

		// create a synthetic block statement, otherwise all hell
		// breaks loose when it comes to block scoping
		raw[ bodyType ] = {
			start: expression.start,
			end: expression.end,
			type: 'BlockStatement',
			body: [ expression ],
			synthetic: true
		};
	}

	Node( raw, parent );

	const type = ( raw.type === 'BlockStatement' ? BlockStatement : types[ raw.type ] ) || Node;
	raw.__proto__ = type.prototype;
}