Error on invalid get component uses

Author: Vorlias

src/Shared/diagnostics.ts

@@ -251,6 +251,28 @@ export const errors = {
 		];
 	}),
 
+	unityMacroExpectsAirshipComponentTypeArgument: errorWithContext((type: string, isUnityObjectType: boolean) => {
+		if (isUnityObjectType) {
+			return [
+				`${type} is a Unity Component, not an Airship Component`,
+				suggestion("Change GetAirshipComponent to GetComponent<" + type + ">()"),
+			];
+		} else {
+			return [`${type} is not a derived type of AirshipBehaviour`];
+		}
+	}),
+
+	unityMacroExpectsComponentTypeArgument: errorWithContext((type: string, isAirshipBehaviourType: boolean) => {
+		if (isAirshipBehaviourType) {
+			return [
+				`${type} is an Airship Component, not a Unity Component`,
+				suggestion("Change GetComponent to GetAirshipComponent<" + type + ">()"),
+			];
+		} else {
+			return [`${type} is not a derived type of Component`];
+		}
+	}),
+
 	decoratorParamsLiteralsOnly: error(
 		"Airship Behaviour decorators only accept literal `string`, `boolean` or `number` values",
 	),
@@ -334,7 +356,7 @@ export const warnings = {
 			"please remove this from the file as it will not be maintained in future.",
 	} satisfies ts.Diagnostic,
 
-	genericBehaviourRefernece: warningWithContext(() => {
+	genericBehaviourReference: warningWithContext(() => {
 		return [
 			"Generic AirshipBehaviours cannot be exposed to the inspector",
 			suggestion(

src/TSTransformer/macros/propertyCallMacros.ts

@@ -4,7 +4,9 @@ import { assert } from "Shared/util/assert";
 import { DiagnosticService } from "TSTransformer/classes/DiagnosticService";
 import { TransformState } from "TSTransformer/classes/TransformState";
 import { MacroList, PropertyCallMacro } from "TSTransformer/macros/types";
+import { isUnityObjectType } from "TSTransformer/util/airshipBehaviourUtils";
 import { convertToIndexableExpression } from "TSTransformer/util/convertToIndexableExpression";
+import { isAirshipBehaviourType, isAirshipBehaviourTypeNode } from "TSTransformer/util/extendsAirshipBehaviour";
 import { isUsedAsStatement } from "TSTransformer/util/isUsedAsStatement";
 import { offset } from "TSTransformer/util/offset";
 import { isDefinitelyType, isNumberType, isStringType } from "TSTransformer/util/types";
@@ -933,6 +935,44 @@ const PROMISE_METHODS: MacroList<PropertyCallMacro> = {
 		}),
 };
 
+const expectAirshipComponentGeneric = (propertyCallMacro: PropertyCallMacro, index: 0 = 0): PropertyCallMacro => {
+	return (state, node, expression, args) => {
+		if (node.typeArguments) {
+			const typeNode = node.typeArguments[index];
+			if (!isAirshipBehaviourTypeNode(state, typeNode)) {
+				DiagnosticService.addDiagnostic(
+					errors.unityMacroExpectsAirshipComponentTypeArgument(
+						node,
+						state.typeChecker.typeToString(state.typeChecker.getTypeFromTypeNode(node.typeArguments[0])),
+						isUnityObjectType(state, state.getType(typeNode)),
+					),
+				);
+			}
+		}
+
+		return propertyCallMacro(state, node, expression, args);
+	};
+};
+
+const expectUnityComponentGeneric = (propertyCallMacro: PropertyCallMacro, index: 0 = 0): PropertyCallMacro => {
+	return (state, node, expression, args) => {
+		if (node.typeArguments) {
+			const type = state.getType(node.typeArguments[index]);
+			if (!isUnityObjectType(state, type)) {
+				DiagnosticService.addDiagnostic(
+					errors.unityMacroExpectsComponentTypeArgument(
+						node,
+						state.typeChecker.typeToString(state.typeChecker.getTypeFromTypeNode(node.typeArguments[0])),
+						isAirshipBehaviourType(state, type, true),
+					),
+				);
+			}
+		}
+
+		return propertyCallMacro(state, node, expression, args);
+	};
+};
+
 const makeTypeArgumentAsStringMacro =
 	(method: string, requiresArgument = true, defaultTypeName?: string): PropertyCallMacro =>
 	(state, node, expression, args) => {
@@ -971,21 +1011,34 @@ const makeTypeArgumentAsStringMacro =
 		}
 	};
 
-const UNITY_GAMEOBJECT_METHODS: MacroList<PropertyCallMacro> = {
+const COMPONENT_MACROS: MacroList<PropertyCallMacro> = {
 	GetComponent: makeTypeArgumentAsStringMacro("GetComponent"),
-	GetAirshipComponent: makeTypeArgumentAsStringMacro("GetAirshipComponent"),
 	GetComponents: makeTypeArgumentAsStringMacro("GetComponents"),
-	GetAirshipComponents: makeTypeArgumentAsStringMacro("GetAirshipComponents"),
 	AddComponent: makeTypeArgumentAsStringMacro("AddComponent"),
+	GetComponentInChildren: makeTypeArgumentAsStringMacro("GetComponentInChildren"),
+	GetComponentsInChildren: makeTypeArgumentAsStringMacro("GetComponentsInChildren"),
 	GetComponentInParent: makeTypeArgumentAsStringMacro("GetComponentInParent"),
+};
+
+const AIRSHIP_COMPONENT_MACROS: MacroList<PropertyCallMacro> = {
+	GetAirshipComponent: makeTypeArgumentAsStringMacro("GetAirshipComponent"),
+	GetAirshipComponents: makeTypeArgumentAsStringMacro("GetAirshipComponents"),
 	AddAirshipComponent: makeTypeArgumentAsStringMacro("AddAirshipComponent"),
-	GetComponentsInChildren: makeTypeArgumentAsStringMacro("GetComponentsInChildren"),
 	GetAirshipComponentsInChildren: makeTypeArgumentAsStringMacro("GetAirshipComponentsInChildren"),
-	GetComponentInChildren: makeTypeArgumentAsStringMacro("GetComponentInChildren"),
 	GetAirshipComponentInChildren: makeTypeArgumentAsStringMacro("GetAirshipComponentInChildren"),
 	GetAirshipComponentInParent: makeTypeArgumentAsStringMacro("GetAirshipComponentInParent"),
 	GetAirshipComponentsInParent: makeTypeArgumentAsStringMacro("GetAirshipComponentsInParent"),
 };
+
+const UNITY_GAMEOBJECT_METHODS: MacroList<PropertyCallMacro> = {};
+
+for (const [macro, call] of Object.entries(AIRSHIP_COMPONENT_MACROS)) {
+	UNITY_GAMEOBJECT_METHODS[macro] = expectAirshipComponentGeneric(call);
+}
+for (const [macro, call] of Object.entries(COMPONENT_MACROS)) {
+	UNITY_GAMEOBJECT_METHODS[macro] = expectUnityComponentGeneric(call);
+}
+
 const UNITY_STATIC_GAMEOBJECT_METHODS: MacroList<PropertyCallMacro> = {
 	FindObjectOfType: makeTypeArgumentAsStringMacro("FindObjectOfType"),
 	FindObjectsByType: makeTypeArgumentAsStringMacro("FindObjectsByType"),

src/TSTransformer/util/airshipBehaviourUtils.ts

@@ -71,6 +71,8 @@ export function getAncestorTypeSymbols(nodeType: ts.Type, typeChecker: ts.TypeCh
 		nodeType = nodeType.getNonNullableType();
 	}
 
+	if (!nodeType.symbol) return [];
+
 	const baseTypes = nodeType.getBaseTypes();
 
 	if (baseTypes) {

src/TSTransformer/util/extendsAirshipBehaviour.ts

@@ -140,7 +140,7 @@ export function isAirshipBehaviourProperty(state: TransformState, node: ts.Prope
 			// Disallow generic properties
 			const nodeTypeRef = node.type;
 			if (nodeTypeRef && ts.isTypeReferenceNode(nodeTypeRef) && nodeTypeRef.typeArguments) {
-				DiagnosticService.addDiagnostic(warnings.genericBehaviourRefernece(node));
+				DiagnosticService.addDiagnostic(warnings.genericBehaviourReference(node));
 				return false;
 			}
 		}
@@ -151,9 +151,11 @@ export function isAirshipBehaviourProperty(state: TransformState, node: ts.Prope
 	return false;
 }
 
-export function isAirshipBehaviourType(state: TransformState, type: ts.Type) {
+export function isAirshipBehaviourType(state: TransformState, type: ts.Type, includeBaseType = false) {
 	const airshipBehaviourSymbol = state.services.airshipSymbolManager.getAirshipBehaviourSymbolOrThrow();
 
+	if (includeBaseType && airshipBehaviourSymbol === type.symbol) return true;
+
 	// Get the inheritance tree, otherwise
 	const inheritance = getAncestorTypeSymbols(type, state.typeChecker);
 	if (inheritance.length === 0) {
@@ -165,6 +167,8 @@ export function isAirshipBehaviourType(state: TransformState, type: ts.Type) {
 	if (baseTypeDeclaration !== undefined) {
 		return baseTypeDeclaration === airshipBehaviourSymbol;
 	}
+
+	return false;
 }
 
 export function isAirshipSingletonType(state: TransformState, type: ts.Type) {
@@ -203,3 +207,15 @@ export function isAirshipSingletonSymbol(state: TransformState, symbol: ts.Symbo
 
 	return false;
 }
+
+export function isAirshipBehaviourTypeNode(state: TransformState, typeNode: ts.TypeNode) {
+	// We'll try the quick symbol way
+	const type = state.getType(typeNode);
+	if (isAirshipBehaviourType(state, type, true)) return true; // quick and dirty
+
+	// If not, we'll use the slower value declaration walking method
+	const symbolOfTypeArg = state.typeChecker.getTypeAtLocation(typeNode).symbol;
+	const valueDeclaration = symbolOfTypeArg.valueDeclaration;
+	if (!valueDeclaration || !ts.isClassDeclaration(valueDeclaration)) return false;
+	return isAirshipBehaviourClass(state, valueDeclaration);
+}