diff --git a/module.json b/module.json
index f139e6f..c54c276 100644
--- a/module.json
+++ b/module.json
@@ -2,7 +2,7 @@
 	"name": "terrain-height-tools",
 	"title": "Terrain Height Tools",
 	"description": "Tools for painting grid cells with terrain heights and calculating line of sight with respect to these heights.",
-	"version": "0.5.1",
+	"version": "0.5.2",
 	"compatibility": {
 		"minimum": "12",
 		"verified": "12",
diff --git a/module/hooks/token-elevation.mjs b/module/hooks/token-elevation.mjs
index 9028245..be5c08a 100644
--- a/module/hooks/token-elevation.mjs
+++ b/module/hooks/token-elevation.mjs
@@ -1,6 +1,6 @@
 import { moduleName, settings } from "../consts.mjs";
 import { TerrainHeightLayer } from "../layers/terrain-height-layer.mjs";
-import { toSceneUnits } from "../utils/grid-utils.mjs";
+import { getSpacesUnderToken, toSceneUnits } from "../utils/grid-utils.mjs";
 import { getTerrainType } from "../utils/terrain-types.mjs";
 
 /**
@@ -61,16 +61,15 @@ export function handleTokenPreCreation(tokenDoc, _createData, _options, userId)
 function getHighestTerrainUnderToken(token, position) {
 	const hm = TerrainHeightLayer.current?._heightMap;
 
-	// We may not want to get the cells under the current position. In this case, we need to work out the offset between
-	// the token's actual position (which is what getOccupiedSpaces returns) and the desired position.
-	const offset = position
-		? { x: position.x - token.x, y: position.y - token.y }
-		: { x: 0, y: 0 };
+	// If a position has been provided, use that position. Otherwise, use the token's position.
+	const { x, y } = position ?? token;
+	const { width, height, hexagonalShape } = token.document;
+	const { type: gridType, size: gridSize } = canvas.grid;
 
 	let highest = 0;
 
-	for (const space of token.getOccupiedSpaces()) {
-		const { i, j } = canvas.grid.getOffset({ x: space.x + offset.x, y: space.y + offset.y });
+	for (const space of getSpacesUnderToken(x, y, width, height, gridType, gridSize, hexagonalShape)) {
+		const { i, j } = canvas.grid.getOffset(space);
 		const terrains = hm.get(i, j);
 		if (!(terrains?.length > 0)) continue; // no terrain at this cell
 
diff --git a/module/utils/grid-utils.mjs b/module/utils/grid-utils.mjs
index 53197d2..10ad573 100644
--- a/module/utils/grid-utils.mjs
+++ b/module/utils/grid-utils.mjs
@@ -1,3 +1,8 @@
+import { cacheReturn } from "./misc-utils.mjs";
+
+/** The side length of a hexagon with a grid size of 1 (apothem of 0.5). */
+const HEX_UNIT_SIDE_LENGTH = 1 / Math.sqrt(3);
+
 /**
  * Returns a set of coordinates for the grid cell at the given position.
  * @param {number} row
@@ -111,3 +116,216 @@ export function fromSceneUnits(val) {
 		? val / canvas.scene.dimensions.distance
 		: null;
 }
+
+const getSquareTokenSpaces = cacheReturn(
+	/**
+	 * Calculates the coordinates of spaces underneath a square token.
+	 * @param {number} width
+	 * @param {number} height
+	 */
+	function(width, height) {
+		/** @type {{ x: number; y: number; }[]} */
+		const spaces = [];
+
+		for (let x = 0; x < width; x++)
+		for (let y = 0; y < height; y++)
+			spaces.push({ x: x + 0.5, y: y + 0.5 });
+
+		return spaces;
+	}
+);
+
+const getEllipseHexTokenSpaces = cacheReturn(
+	/**
+	 * Calculates the cube coordinates of all spaces occupied by an ellipse token with the given width/height.
+	 * @param {number} primaryAxisSize Size of the token in the primary direction, measured in cells.
+	 * @param {number} secondaryAxisSize Size of the token in the secondary direction, measured in cells.
+	 * @param {boolean} isColumnar true for hex columns, false for hex rows.
+	 * @param {boolean} isVariant2 false for ELLIPSE_1, true for ELLIPSE_2.
+	 */
+	function(primaryAxisSize, secondaryAxisSize, isColumnar, isVariant2) {
+		// Ellipses require the size in primary axis to be at least as big as `floor(secondaryAxisSize / 2) + 1`.
+		// E.G. for columnar grids, for a width of 5, the height must be 3 or higher. For a width of 6, height must be
+		// at least 4 or higher. Same is true for rows, but in the opposite axis.
+		if (primaryAxisSize < Math.floor(secondaryAxisSize / 2) + 1) {
+			return [];
+		}
+
+		const secondaryAxisOffset = Math[isVariant2 ? "ceil" : "floor"]((secondaryAxisSize - 1) / 2) * HEX_UNIT_SIDE_LENGTH * 1.5 + HEX_UNIT_SIDE_LENGTH;
+
+		/** @type {{ x: number; y: number; }[]} */
+		const spaces = [];
+
+		// Track the offset distance from the largest part of the hex (in primary), and which side we're on.
+		// The initial side we use (sign) depends on the variant of ellipse we're building.
+		let offsetDist = 0;
+		let offsetSign = isVariant2 ? 1 : -1;
+
+		for (let i = 0; i < secondaryAxisSize; i++) {
+			const primaryAxisOffset = (offsetDist + 1) / 2;
+			const secondaryPosition = offsetDist * offsetSign * HEX_UNIT_SIDE_LENGTH * 1.5 + secondaryAxisOffset;
+
+			// The number of spaces in this primary axis decreases by 1 each time the offsetDist increases by 1: at the
+			// 0 (the largest part of the shape), we have the full primary size number of cells. Either side of this, we
+			// have primary - 1, either side of those primary - 2, etc.
+			for (let j = 0; j < primaryAxisSize - offsetDist; j++) {
+				spaces.push(coordinate(j + primaryAxisOffset, secondaryPosition));
+			}
+
+			// Swap over the offset side, and increase dist if neccessary
+			offsetSign *= -1;
+			if (i % 2 === 0) offsetDist++;
+		}
+
+		return spaces;
+
+		/**
+		 * @param {number} primary
+		 * @param {number} secondary
+		 */
+		function coordinate(primary, secondary) {
+			return isColumnar ? { x: secondary, y: primary } : { x: primary, y: secondary };
+		}
+	}
+);
+
+const getTrapezoidHexTokenSpaces = cacheReturn(
+	/**
+	 * Calculates the cube coordinates of all spaces occupied by an trapezoid token with the given width/height.
+	 * @param {number} primaryAxisSize Size of the token in the primary direction, measured in cells.
+	 * @param {number} secondaryAxisSize Size of the token in the secondary direction, measured in cells.
+	 * @param {boolean} isColumnar true for hex columns, false for hex rows.
+	 * @param {boolean} isVariant2 false for TRAPEZOID_1, true for TRAPEZOID_2.
+	 */
+	function(primaryAxisSize, secondaryAxisSize, isColumnar, isVariant2) {
+		// For trapezoid to work, the size in the primary axis must be equal to or larger than the size in the secondary
+		if (primaryAxisSize < secondaryAxisSize) {
+			return [];
+		}
+
+		const secondaryAxisOffset = isVariant2 ? HEX_UNIT_SIDE_LENGTH + (secondaryAxisSize - 1) * HEX_UNIT_SIDE_LENGTH * 1.5 : HEX_UNIT_SIDE_LENGTH;
+
+		/** @type {{ x: number; y: number; }[]} */
+		const spaces = [];
+
+		// Trazpezoids are simple. Start with a line in the primary direction that is the full primary size.
+		// Then, for each cell in the secondary direction, reduce the primary by one.
+		// If we are doing variant1 we offset in the secondary by one direction, for variant2 we go the other direction.
+		for (let i = 0; i < secondaryAxisSize; i++) {
+			const primaryAxisOffset = (i + 1) / 2;
+			const secondaryPosition = i * (isVariant2 ? -1 : 1) * HEX_UNIT_SIDE_LENGTH * 1.5 + secondaryAxisOffset;
+
+			for (let j = 0; j < primaryAxisSize - i; j++) {
+				spaces.push(coordinate(j + primaryAxisOffset, secondaryPosition));
+			}
+		}
+
+		return spaces;
+
+		/**
+		 * @param {number} primary
+		 * @param {number} secondary
+		 */
+		function coordinate(primary, secondary) {
+			return isColumnar ? { x: secondary, y: primary } : { x: primary, y: secondary };
+		}
+	}
+);
+
+const getRectangleHexTokenSpaces = cacheReturn(
+	/**
+	 * Calculates the cube coordinates of all spaces occupied by an trapezoid token with the given width/height.
+	 * @param {number} primaryAxisSize Size of the token in the primary direction, measured in cells.
+	 * @param {number} secondaryAxisSize Size of the token in the secondary direction, measured in cells.
+	 * @param {boolean} isColumnar true for hex columns, false for hex rows.
+	 * @param {boolean} isVariant2 false for TRAPEZOID_1, true for TRAPEZOID_2.
+	 */
+	function(primaryAxisSize, secondaryAxisSize, isColumnar, isVariant2) {
+		// If the size in the primary direction is 1, the size in the secondary direction must be no more than one.
+		// For primary size >= 2, any size secondary is acceptable.
+		if (primaryAxisSize === 1 && secondaryAxisOffset > 1) {
+			return [];
+		}
+
+		/** @param {{ x: number; y: number; }[]} */
+		const spaces = [];
+
+		const largeRemainder = isVariant2 ? 1 : 0;
+
+		// Spaces under rectangles are easy. They just alternate size in the primary direction by 0 and -1 as we iterate
+		// through the cells in the secondary direction.
+		for (let i = 0; i < secondaryAxisSize; i++) {
+			const isLarge = i % 2 === largeRemainder;
+			for (let j = 0; j < primaryAxisSize - (isLarge ? 0 : 1); j++) {
+				spaces.push(coordinate(
+					j + (isLarge ? 0.5 : 1),
+					i * HEX_UNIT_SIDE_LENGTH * 1.5 + HEX_UNIT_SIDE_LENGTH));
+			}
+		}
+
+		return spaces;
+
+		/**
+		 * @param {number} primary
+		 * @param {number} secondary
+		 */
+		function coordinate(primary, secondary) {
+			return isColumnar ? { x: secondary, y: primary } : { x: primary, y: secondary };
+		}
+	}
+);
+
+/**
+ * @param {number} x Token X.
+ * @param {number} y Token Y.
+ * @param {number} width Token width (in grid spaces).
+ * @param {number} height Token height (in grid spaces).
+ * @param {number} gridType The type of grid.
+ * @param {number} gridSize The size of the grid in pixels.
+ * @param {number} hexShape For hexagonal tokens, the type of hex shape used.
+ */
+export function getSpacesUnderToken(x, y, width, height, gridType, gridSize, hexShape) {
+	// Gridless is not supported
+	if (gridType === CONST.GRID_TYPES.GRIDLESS) {
+		return [];
+	}
+
+	// For square, can easily work the points out by enumerating over the width/height
+	if (gridType === CONST.GRID_TYPES.SQUARE) {
+		return getSquareTokenSpaces(width, height)
+			.map(p => ({ x: x + p.x * gridSize, y: y + p.y * gridSize }));
+	}
+
+	// For hex grids, it depends on the token's hex shape:
+	// Hex grids are also rotationally equivalent (i.e. for a hex row we can just swap X and Y from a hex column).
+	// We define a "primary" axis (the direction in the namesake of the grid - i.e. Y/height for columns and X/width for
+	// rows). The "secondary" axis is the other (X/height for columns, Y/height for rows).
+	const isColumnar = [CONST.GRID_TYPES.HEXEVENQ, CONST.GRID_TYPES.HEXODDQ].includes(gridType);
+	const primaryAxisSize = isColumnar ? height : width;
+	const secondaryAxisSize = isColumnar ? width : height;
+	const isVariant2 = [
+		CONST.TOKEN_HEXAGONAL_SHAPES.ELLIPSE_2,
+		CONST.TOKEN_HEXAGONAL_SHAPES.TRAPEZOID_2,
+		CONST.TOKEN_HEXAGONAL_SHAPES.RECTANGLE_2
+	].includes(hexShape);
+
+	switch (hexShape) {
+		case CONST.TOKEN_HEXAGONAL_SHAPES.ELLIPSE_1:
+		case CONST.TOKEN_HEXAGONAL_SHAPES.ELLIPSE_2:
+			return getEllipseHexTokenSpaces(primaryAxisSize, secondaryAxisSize, isColumnar, isVariant2)
+				.map(p => ({ x: x + p.x * gridSize, y: y + p.y * gridSize }));
+
+		case CONST.TOKEN_HEXAGONAL_SHAPES.TRAPEZOID_1:
+		case CONST.TOKEN_HEXAGONAL_SHAPES.TRAPEZOID_2:
+			return getTrapezoidHexTokenSpaces(primaryAxisSize, secondaryAxisSize, isColumnar, isVariant2)
+				.map(p => ({ x: x + p.x * gridSize, y: y + p.y * gridSize }));
+
+		case CONST.TOKEN_HEXAGONAL_SHAPES.RECTANGLE_1:
+		case CONST.TOKEN_HEXAGONAL_SHAPES.RECTANGLE_2:
+			return getRectangleHexTokenSpaces(primaryAxisSize, secondaryAxisSize, isColumnar, isVariant2)
+				.map(p => ({ x: x + p.x * gridSize, y: y + p.y * gridSize }));
+
+		default:
+			throw new Error("Unknown hex grid type.");
+	}
+}
diff --git a/module/utils/misc-utils.mjs b/module/utils/misc-utils.mjs
index f9d7bb0..c589320 100644
--- a/module/utils/misc-utils.mjs
+++ b/module/utils/misc-utils.mjs
@@ -88,3 +88,24 @@ export class OrderedSet {
 		};
 	}
 }
+
+/**
+ * Wraps a function such that the return value of the function is cached based on the parameters passed to it.
+ * @template {(...args: any) => any} T
+ * @param {T} func Function to wrap.
+ * @param {(args: Parameters<T>) => string} keyFunc Optional function to generate cache key from the arguments. If not
+ * provided, defaults to all arguments joined by a "|".
+ * @returns {T}
+ */
+export function cacheReturn(func, keyFunc = undefined) {
+	const cache = new Map();
+	keyFunc ??= args => args.join("|");
+
+	return function(...args) {
+		const key = keyFunc(args);
+		if (cache.has(key)) return cache.get(key);
+		const result = func(...args);
+		cache.set(key, result);
+		return result;
+	}
+}