Macro to generated structured field access to a raw byte array

StructView.hx

/**
 * StructView – Typed access of a byte buffer
 * 
 * Example:
 * ```
 * typedef MessageView = StructView<{
 * 	a: Int,
 * 	b: Float,
 * 	sub: {
 * 		c: StructView.UInt8,
 * 		d: Float
 * 	}
 * }>;
 * ```
 * 
 * Bytes are tightly packed and match the order defined in the structure
 * 
 * @author George Corney (haxiomic)
 * @license MIT
 * @source https://gist.github.com/haxiomic/11efc43b01792718deb3d68aa901a457
 **/

#if !macro

@:genericBuild(StructView.buildModule())
class StructView<T> {
	public function new() {} // stub to help autocomplete
}

typedef ArrayBuffer = js.lib.ArrayBuffer;
typedef DataView = js.lib.DataView;

// fixed size data types
@:coreType abstract Float64 from Float to Float {}
@:coreType abstract Float32 from Float to Float {}
@:coreType abstract Int32 from Int to Int {}
@:coreType abstract Int16 from Int to Int {}
@:coreType abstract Int8 from Int to Int {}
@:coreType abstract UInt32 from Int to Int {}
@:coreType abstract UInt16 from Int to Int {}
@:coreType abstract UInt8 from Int to Int {}

#else

import haxe.macro.Context;
import haxe.macro.TypeTools;
import haxe.macro.Expr;

@:persistent var generatedTypeIndex = 0;

function buildModule() {
	// extract name and structure from type parameters
	var typeParams = switch Context.followWithAbstracts(Context.getLocalType()) {
		case TInst(_, [
			anon = Context.followWithAbstracts(_) => TAnonymous(a)
		]):
			{
				anon: anon,
				fields: a.get().fields,
				name: 'StructView' + generatedTypeIndex++,
				byteOffset: 0,
			};
		case type:
			Context.fatalError('Type parameter must be a structure ($type)', Context.currentPos());
	}

	var moduleName = typeParams.name;
	var moduleTypes = getModuleTypes(typeParams);

	Context.defineModule('structview.$moduleName', moduleTypes);
	
	return macro : structview.$moduleName;
}

function getModuleTypes(
	params: {
		name: String,
		anon: haxe.macro.Type,
		fields: Array<haxe.macro.Type.ClassField>,
		byteOffset: Int,
	},
	?parentTypeName
) {
	var subModules = [];

	var generatedTypeName = params.name;

	// define abstract type
	var anonTypeComplex = TypeTools.toComplexType(params.anon);

	// determine byte length
	var byteLengths = [for (field in params.fields)
		getFieldByteLength(field)
	];

	var totalByteLength = 0;
	for (x in byteLengths) totalByteLength += x;

	var generatedType = macro class $generatedTypeName { };
	generatedType.kind = TDAbstract(macro: StructView.DataView);
	generatedType.doc = 'Typed access to byte buffer (macro-generated type)\n\n' + generateMarkdownTable(params.fields, byteLengths) + '\nTotal bytes: $totalByteLength';
	generatedType.meta = [
		{
			name: ':forward',
			params: [
				macro byteLength,
				macro byteOffset,
				macro buffer
			],
			pos: generatedType.pos
		}
	];
	
	// add new()
	if (parentTypeName == null) {
		var newFields = (macro class {
			public inline function new(?arrayBuffer: StructView.ArrayBuffer, byteOffset: Int = 0, ?fields: $anonTypeComplex) {
				if (arrayBuffer == null) {
					arrayBuffer = new StructView.ArrayBuffer($v{totalByteLength});
					byteOffset = 0;
				}
				this = new StructView.DataView(arrayBuffer, byteOffset);
				if (fields != null) {
					set(fields);
				}
			}
		}).fields;

		for (f in newFields) generatedType.fields.push(f);
	} else {
		var newFields = (macro class {
			public inline function new(data: StructView.DataView) {
				this = data;
			}
		}).fields;
		for (f in newFields) generatedType.fields.push(f);
	}

	// add byteLength
	generatedType.fields.push((macro class {
		static public final BYTE_LENGTH = $v{totalByteLength};
	}).fields[0]);

	// add getter and setter fields
	for (i => field in params.fields) {
		var name = field.name;

		var subStructField = false;

		var byteOffset = {
			var o = params.byteOffset;
			for (j in 0...i) o += byteLengths[j];
			o;
		}

		var complexType: ComplexType = switch Context.followWithAbstracts(field.type) {
			case anon = TAnonymous(a):
				var subTypeName =
					'${generatedTypeName}_' +
					name.substr(0, 1).toUpperCase() + name.substr(1);
				
				// we need to build a sub type for this field
				subModules = subModules.concat(
					getModuleTypes(
						{
							name: subTypeName,
							fields: a.get().fields,
							anon: anon,
							byteOffset: byteOffset,
						},
						generatedTypeName
					)
				);

				subStructField = true;

				TPath({name: subTypeName, pack: []});
			default: TypeTools.toComplexType(field.type);
		}

		var bufferName = 'this';
		var byteOffsetExpr = macro $v{byteOffset};
 		
		var get_name = 'get_$name';
		var set_name = 'set_$name';
		var newFields = if (subStructField) {
			var subTypePath = switch complexType {
				case TPath(p):p;
				default: throw 'Expected TPath';
			}
			(macro class {
				public var $name(get, never): $complexType;
				inline function $get_name(): $complexType {
					return new $subTypePath(this);
				}
			}).fields;
		} else {
		 	(macro class {
				public var $name(get, set): $complexType;
				inline function $get_name(): $complexType {
					return ${getReadExpr(field, bufferName, byteOffsetExpr)}
				}
				inline function $set_name(v: $complexType) {
					${getWriteExpr(field, bufferName, byteOffsetExpr)}
					return v;
				}
			}).fields;
		}

		for (newField in newFields) {
			generatedType.fields.push(newField);
		}
	}

	// add set(obj)
	generatedType.fields.push({
		var setExpr = [for (field in params.fields) {
			var name = field.name;

			switch Context.followWithAbstracts(field.type) {
				case anon = TAnonymous(a):
					macro $i{name}.set(values.$name);
				default:
					macro $i{name} = values.$name;
			}
		}];
		(macro class {
			public inline function set(values: $anonTypeComplex) {
				$b{setExpr};
			}
		}).fields[0];
	});

	// add toString()
	generatedType.fields.push({
		var lineExprs = [for (field in params.fields) {
			var name = field.name;
			switch Context.followWithAbstracts(field.type) {
				case anon = TAnonymous(a):
					macro str += '\n$tabDepth' + $v{name} + ': ' + $i{name}.toString(tabDepth + '\t');
				default:
					macro str += '\n$tabDepth' + $v{name} + ': ' + $i{name};
			}
		}];
		(macro class {
			public function toString(?tabDepth = '\t'): String {
				var str = '';
				var name = $v{generatedTypeName};
				str += '$name [$this] {';
				$b{lineExprs}
				str += '\n${tabDepth.substr(1)}}';
				return str;
			}
		}).fields[0];
	});

	// debug generated types:
	// trace(new haxe.macro.Printer().printTypeDefinition(generatedType, false));

	return [generatedType].concat(subModules);
}

function getFieldByteLength(field: haxe.macro.Type.ClassField): Int {
	var resolved = Context.followWithAbstracts(field.type);
	var byteLength = switch resolved {
		case TAbstract(_.get() => t, []):
			switch t {
				case {module: 'StdTypes', name: 'Float'}: 8;
				case {module: 'StdTypes', name: 'Int'}: 4;
				case {module: 'StdTypes', name: 'Bool'}: 1;

				case {module: 'StructView', name: 'Float64'}: 8;
				case {module: 'StructView', name: 'Float32'}: 4;
				case {module: 'StructView', name: 'Int32'}: 4;
				case {module: 'StructView', name: 'Int16'}: 2;
				case {module: 'StructView', name: 'Int8'}: 1;
				case {module: 'StructView', name: 'UInt32'}: 4;
				case {module: 'StructView', name: 'UInt16'}: 2;
				case {module: 'StructView', name: 'UInt8'}: 1;
				
				default: null;
			}
		case TInst(_.get() => t, []):
			switch t {
				case {module: 'haxe.Int64', name: '___Int64'}: 8;
				default: null;
			}
		case TAnonymous(_.get() => anon):
			var structLength = 0;
			for (f in anon.fields) {
				structLength += getFieldByteLength(f);
			}
			structLength;
		default:
			null;
	}

	if (byteLength == null) {
		Context.error('StructView.getFieldByteLength: Unsupported type ${field.type}', field.pos);
	}

	return byteLength;
}

function getReadExpr(field: haxe.macro.Type.ClassField, buffer: String, byteOffsetExpr: Expr) {
	var expr = switch Context.followWithAbstracts(field.type) {
		case TAbstract(_.get() => t, []):
			switch t {
				case {module: 'StdTypes', name: 'Float'}: macro $i{buffer}.getFloat64($byteOffsetExpr);
				case {module: 'StdTypes', name: 'Int'}: macro cast $i{buffer}.getInt32($byteOffsetExpr);
				case {module: 'StdTypes', name: 'Bool'}: macro cast $i{buffer}.get($byteOffsetExpr);

				case {module: 'StructView', name: 'Float64'}: macro cast $i{buffer}.getFloat64($byteOffsetExpr);
				case {module: 'StructView', name: 'Float32'}: macro cast $i{buffer}.getFloat32($byteOffsetExpr);
				case {module: 'StructView', name: 'Int32'}: macro cast $i{buffer}.getInt32($byteOffsetExpr);
				case {module: 'StructView', name: 'Int16'}: macro cast $i{buffer}.getInt16($byteOffsetExpr);
				case {module: 'StructView', name: 'Int8'}: macro cast $i{buffer}.getInt8($byteOffsetExpr);
				case {module: 'StructView', name: 'UInt32'}: macro cast $i{buffer}.getUint32($byteOffsetExpr);
				case {module: 'StructView', name: 'UInt16'}: macro cast $i{buffer}.getUint16($byteOffsetExpr);
				case {module: 'StructView', name: 'UInt8'}: macro cast $i{buffer}.getUint8($byteOffsetExpr);

				default:
					null;
			}
		case TInst(_.get() => t, []):
			switch t {
				case {module: 'haxe.Int64', name: '___Int64'}: macro $i{buffer}.getInt64($byteOffsetExpr);
				default:
					null;
			}
		case TAnonymous(_.get() => anon): macro null;
		case t:
			null;
	}

	if (expr == null) {
		Context.error('StructView.getReadExpr: Unsupported type ${field.type}', field.pos);
	}

	return expr;
}

function getWriteExpr(field: haxe.macro.Type.ClassField, buffer: String, byteOffsetExpr: Expr) {
	var resolved = Context.followWithAbstracts(field.type);
	var expr = switch resolved {
		case TAbstract(_.get() => t, []):
			switch t {
				case {module: 'StdTypes', name: 'Float'}: macro $i{buffer}.setFloat64($byteOffsetExpr, v);
				case {module: 'StdTypes', name: 'Int'}: macro $i{buffer}.setInt32($byteOffsetExpr, cast v);
				case {module: 'StdTypes', name: 'Bool'}: macro cast $i{buffer}.set($byteOffsetExpr, v ? 1 : 0);

				case {module: 'StructView', name: 'Float64'}: macro $i{buffer}.setFloat64($byteOffsetExpr, v);
				case {module: 'StructView', name: 'Float32'}: macro $i{buffer}.setFloat32($byteOffsetExpr, v);
				case {module: 'StructView', name: 'Int32'}: macro $i{buffer}.setInt32($byteOffsetExpr, v);
				case {module: 'StructView', name: 'Int16'}: macro $i{buffer}.setInt16($byteOffsetExpr, v);
				case {module: 'StructView', name: 'Int8'}: macro $i{buffer}.setInt8($byteOffsetExpr, v);
				case {module: 'StructView', name: 'UInt32'}: macro $i{buffer}.setUint32($byteOffsetExpr, v);
				case {module: 'StructView', name: 'UInt16'}: macro $i{buffer}.setUint16($byteOffsetExpr, v);
				case {module: 'StructView', name: 'UInt8'}: macro $i{buffer}.setUint8($byteOffsetExpr, v);

				default: null;
			}
		case TInst(_.get() => t, []):
			switch t {
				case {module: 'haxe.Int64', name: '___Int64'}: macro $i{buffer}.setInt64($byteOffsetExpr, v);
				default: null;
			}
		case TAnonymous(_.get() => anon): macro null;
		default:
			null;
	}

	if (expr == null) {
		Context.error('StructView.getWriteExpr: Unsupported type ${field.type}', field.pos);
	}

	return expr;
}

function generateMarkdownTable(fields: Array<haxe.macro.Type.ClassField>, byteLengths: Array<Int>) {
	var table = '| Name | Type | Byte Offset | Byte Length |\n';
	table += '| --- | --- | --- | --- |\n';
	var byteOffset = 0;
	for (i in 0...fields.length) {
		var field = fields[i];
		var byteLength = byteLengths[i];
		table += '| ${field.name} | ${TypeTools.toString(field.type)} | $byteOffset | $byteLength |\n';
		byteOffset += byteLength;
	}
	return table;
}

#end