ShortLambda.hx

package ;

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

using Lambda;

class ShortLambda {
	public static function build(): Array<Field> {
		var fields = Context.getBuildFields();

		for (field in fields) {
			switch (field.kind) {
			case FFun(fun): 
				switch (fun.expr.expr) {
				case EBlock(blocks):
					for (b in blocks) {
						switch (b.expr) {
						case EVars([{ 
							expr: { expr: EMeta({ name:":fun", params:args, pos:pos }, expr), pos:_ },
							name: name,
							type: type
						}]):
							b.expr = EVars([{ 
								expr: processLambdaClosure(args, expr, pos),
								name: name, type: type 
							}]);
						default:
						}
					}
				default:
				}
			default:
			}
		}

		return fields;
	}

	
	private static function processLambdaClosure(args: Array<Expr>, expr: Expr, p: Position): Expr {
		var funArgs: Array<FunctionArg> = 
			args
			.mapi(function(i, arg: Expr) {
				var name = switch (arg.expr) {
				case EConst(CIdent(n)):
					n;
				default:
					'_${i}';
				}
				return { name: name, opt: false, type: null, value: null };
			})
			.array()
		;

		var blocksExpr = 
			switch (expr.expr) {
			case EBlock(blocks): { expr: EBlock(processBlocks(blocks)), pos: expr.pos };
			default: expr;
			}

		return {
            expr: EFunction(null, {
                args: funArgs, 
                params: [], 
                ret: null,
                expr: blocksExpr
            }),
            pos: p
        };
	}

	private static function processBlocks(blocks: Array<Expr>): Array<Expr> {
		if (blocks.length > 0) {
			var last = blocks.pop();

			switch (last.expr) {
			case EReturn(_), EVars(_): blocks.push(last);
			default: blocks.push(asReturn(last));
			}
		}

		return blocks;
	}

	private static function asReturn(expr: Expr): Expr {
		return { 
			expr: EReturn(expr),
			pos: Context.currentPos()
		};
	}
}