WIP 2

TODO.md

@@ -2,8 +2,7 @@
 
 * Prevent injection of class into scope function for its methods.
 
-* Ensure __proto__ of class and class prototype are always set correctly.
-  * Tests for `class extends Object {}` and `class extends Function {}`
+* Tests for `class extends Object {}` and `class extends Function {}`
 
 * Tests for `super()` in arrow functions:
 

lib/serialize/functions.js

@@ -262,7 +262,7 @@ module.exports = {
 		// Return node with placeholder for function definition - definition will be added later
 		return this.wrapFunctionWithProperties(
 			fn, record, t.identifier('x'), fnDef.name, fnDef.numParams,
-			isClass, isAsync, isGenerator, fnDef.isArrow, fnDef.isMethod
+			isClass, fnDef.isClassWithSuperClass, isAsync, isGenerator, fnDef.isArrow, fnDef.isMethod
 		);
 	},
 
@@ -462,7 +462,7 @@ module.exports = {
 		// fn, record, node, name, numParams, isClass, isAsync, isGenerator, isArrowOrBound, isMethod
 		return this.wrapFunctionWithProperties(
 			fn, record, node, isCircular ? '' : `bound ${unboundFn.name}`, numParams,
-			false, false, false, true, false
+			false, false, false, false, true, false
 		);
 	},
 
@@ -526,7 +526,8 @@ module.exports = {
 	},
 
 	wrapFunctionWithProperties(
-		fn, record, node, name, numParams, isClass, isAsync, isGenerator, isArrowOrBound, isMethod
+		fn, record, node, name, numParams,
+		isClass, isClassWithSuperClass, isAsync, isGenerator, isArrowOrBound, isMethod
 	) {
 		// Set default `length` + `name` properties based on function definition
 		const defaultProps = [
@@ -537,7 +538,11 @@ module.exports = {
 		// Get `prototype` property and set default `prototype` prop
 		const defaultProto = isGenerator
 			? isAsync ? AsyncGeneratorPrototype : GeneratorPrototype
-			: isAsync ? AsyncFunctionPrototype : Function.prototype;
+			: isAsync
+				? AsyncFunctionPrototype
+				: isClassWithSuperClass
+					? undefined // Has been defined with `extends class {}`. Always need to set prototype.
+					: Function.prototype;
 
 		const protoDescriptor = Object.getOwnPropertyDescriptor(fn, 'prototype'),
 			proto = protoDescriptor ? protoDescriptor.value : undefined;
@@ -604,7 +609,7 @@ module.exports = {
 				protoIsAltered = true;
 			} else if (propNames[0] !== 'constructor') {
 				protoIsAltered = true;
-			} else if (Object.getPrototypeOf(proto) !== Object.prototype) {
+			} else if (isClassWithSuperClass || Object.getPrototypeOf(proto) !== Object.prototype) {
 				protoIsAltered = true;
 			} else {
 				const ctorDescriptor = Object.getOwnPropertyDescriptor(proto, 'constructor');
@@ -615,9 +620,11 @@ module.exports = {
 
 			if (protoIsAltered) {
 				const protoRecord = this.serializeValue(proto, this.getPrototypeVarName(record), '.prototype');
+				// Classes with super class are defined with `extends class {}`, so always need to set prototype
+				const defaultProtoProto = isClassWithSuperClass ? undefined : Object.prototype;
 				this.withTrace(
 					() => this.serializeProperties(
-						proto, protoRecord, null, Object.prototype,
+						proto, protoRecord, null, defaultProtoProto,
 						[{name: 'constructor', value: fn, writable: true, enumerable: false, configurable: true}],
 						undefined, undefined, true
 					),

lib/serialize/parseFunction.js

@@ -51,6 +51,7 @@ const RUNTIME_DIR_PATH = pathJoin(__dirname, '../runtime/');
  *   {string} .name - `.name` of created function
  *   {number} .numParams - `.length` of created function
  *   {boolean} .isClass - `true` if is class
+ *   {boolean} .isClassWithSuperClass - `true` if is class with super class (`extends ...`)
  *   {boolean} .isAsync - `true` if is async function
  *   {boolean} .isGenerator - `true` if is generator
  *   {boolean} .isArrow - `true` if is arrow function
@@ -163,7 +164,8 @@ module.exports = function parseFunction(
 		name;
 	const {type} = node,
 		isArrow = type === 'ArrowFunctionExpression',
-		containsEval = !!fnInfo.containsEval;
+		containsEval = !!fnInfo.containsEval,
+		isClassWithSuperClass = !!fnInfo.hasSuperClass;
 	if (isArrow) {
 		// Arrow function
 		name = '';
@@ -265,7 +267,7 @@ module.exports = function parseFunction(
 
 				// Replace `extends` clause with `extends class {}`
 				// TODO: Make sure prototypes get changed in all cases.
-				if (fnInfo.hasSuperClass) node.superClass = t.classExpression(null, null, t.classBody([]));
+				if (isClassWithSuperClass) node.superClass = t.classExpression(null, null, t.classBody([]));
 			}
 		}
 	}
@@ -312,6 +314,7 @@ module.exports = function parseFunction(
 		name,
 		numParams,
 		isClass,
+		isClassWithSuperClass,
 		isAsync,
 		isGenerator,
 		isArrow,

test/classes.test.js

@@ -3358,7 +3358,8 @@ describe('Classes', () => {
 			},
 			out: `(()=>{
 				const a=Object,
-					b=(b=>[
+					b=a.setPrototypeOf,
+					c=(b=>[
 						class Klass extends class{}{
 							constructor(){
 								super();
@@ -3375,14 +3376,17 @@ describe('Classes', () => {
 							}.foo
 						]
 					])(1),
-					c=a.setPrototypeOf(b[0],a),
-					d=b[1]();
-				d[0](c);
-				a.defineProperties(
-					c.prototype,
-					{foo:{value:d[1],writable:true,configurable:true}}
+					d=b(c[0],a),
+					e=c[1]();
+				e[0](d);
+				b(
+					a.defineProperties(
+						d.prototype,
+						{foo:{value:e[1],writable:true,configurable:true}}
+					),
+					a.prototype
 				);
-				return c
+				return d
 			})()`,
 			validate(Klass) {
 				expect(Klass).toBeFunction();