shader link working

hxsl/Linker.hx

@@ -1,23 +1,43 @@
 package hxsl;
 using hxsl.Ast;
 
+private enum VarKind {
+	Var;
+	Function;
+	Vertex;
+	Fragment;
+	Init;
+}
+
 private class AllocatedVar {
 	public var id : Int;
 	public var v : TVar;
 	public var path : String;
 	public var merged : Array<TVar>;
+	public var kind : VarKind;
+	public var parent : AllocatedVar;
 	public function new() {
+		kind = Var;
 	}
 }
 
 private class ShaderInfos {
+	public var name : String;
 	public var priority : Int;
 	public var body : TExpr;
-	public var deps : Array<TFunction>;
+	public var usedFunctions : Array<TFunction>;
+	public var deps : Map<ShaderInfos, Bool>;
 	public var read : Map<Int,AllocatedVar>;
 	public var write : Map<Int,AllocatedVar>;
-	public function new() {
-		deps = [];
+	public var processed : Map<Int, Bool>;
+	public var vertex : Bool;
+	public var onStack : Bool;
+	public var marked : Null<Bool>;
+	public function new(n, v) {
+		this.name = n;
+		this.vertex = v;
+		processed = new Map();
+		usedFunctions = [];
 		read = new Map();
 		write = new Map();
 	}
@@ -27,6 +47,7 @@ class Linker {
 
 	var varMap : Map<String,AllocatedVar>;
 	var allVars : Array<AllocatedVar>;
+	var allFuns : Map<Int,TFunction>;
 	var curShader : ShaderInfos;
 	var shaders : Array<ShaderInfos>;
 	
@@ -67,11 +88,11 @@ class Linker {
 		}
 	}
 	
-	function allocVar( v : TVar, ?path : String, ?parent : TVar, p : Position ) : AllocatedVar {
+	function allocVar( v : TVar, p : Position, ?path : String, ?parent : TVar ) : AllocatedVar {
 		if( v.kind == Local )
 			throw "assert";
 		if( v.parent != null && parent == null ) {
-			parent = allocVar(v.parent, null, null, p).v;
+			parent = allocVar(v.parent, p).v;
 			var p = parent;
 			path = p.name;
 			p = p.parent;
@@ -81,6 +102,12 @@ class Linker {
 			}
 		}
 		var key = (path == null ? v.name : path + "." + v.name);
+		if( v.qualifiers != null )
+			for( q in v.qualifiers )
+				switch( q ) {
+				case Name(n): key = n;
+				default:
+				}
 		var v2 = varMap.get(key);
 		if( v2 != null ) {
 			for( vm in v2.merged )
@@ -111,11 +138,12 @@ class Linker {
 		a.merged = [v];
 		a.path = key;
 		a.id = allVars.length;
+		a.parent = parent == null ? null : allocVar(parent, p);
 		allVars.push(a);
 		varMap.set(key, a);
 		switch( v2.type ) {
 		case TStruct(vl):
-			v2.type = TStruct([for( v in vl ) allocVar(v, key, v2, p).v]);
+			v2.type = TStruct([for( v in vl ) allocVar(v, p, key, v2).v]);
 		default:
 		}
 		return a;
@@ -147,26 +175,199 @@ class Linker {
 	function mapExprVar( e : TExpr ) {
 		switch( e.e ) {
 		case TVar(v) if( v.kind != Local ):
-			var v = allocVar(v, null, null, e.p);
+			var v = allocVar(v, e.p);
+			if( curShader != null ) {
+				//trace(curShader.name + " read " + v.path);
+				curShader.read.set(v.id, v);
+			}
 			return { e : TVar(v.v), t : v.v.type, p : e.p };
+		case TBinop(op, e1, e2):
+			switch( [op,e1.e] ) {
+			case [OpAssign | OpAssignOp(_), (TVar(v) | TSwiz({ e : TVar(v) },_))] if( v.kind != Local ):
+				var v = allocVar(v, e1.p);
+				if( curShader != null ) {
+					//trace(curShader.name + " write " + v.path);
+					curShader.write.set(v.id, v);
+				}
+				// TSwiz might only assign part of the components, let's then consider that we read the other
+				if( op == OpAssign && (switch( e1.e ) { case TVar(_): true; default: false; }) ) {
+					var eout = { e : TVar(v.v), t : e1.t, p : e1.p };
+					return { e : TBinop(OpAssign, eout, mapExprVar(e2)), t : e.t, p : e.p };
+				}
+			default:
+			}
 		default:
-			return e.map(mapExprVar);
 		}
+		return e.map(mapExprVar);
+	}
+	
+	function addShader( name : String, vertex : Bool, e : TExpr, p : Int ) {
+		curShader = new ShaderInfos(name, vertex);
+		curShader.priority = p;
+		curShader.body = mapExprVar(e);
+		shaders.push(curShader);
+		curShader = null;
+	}
+	
+	function sortByPriorityDesc( s1 : ShaderInfos, s2 : ShaderInfos ) {
+		return s2.priority - s1.priority;
+	}
+	
+	function buildDependency( parent : ShaderInfos, v : AllocatedVar ) {
+		for( s in shaders ) {
+			if( !s.write.exists(v.id) || parent == s )
+				continue;
+			//trace(parent.name + " => " + s.name + " (" + v.path + ")");
+			parent.deps.set(s, true);
+			if( s.deps == null ) {
+				s.deps = new Map();
+				for( r in s.read )
+					buildDependency(s, r);
+			}
+			if( !s.read.exists(v.id) )
+				return;
+		}
+		if( v.v.kind == Var )
+			error("Variable " + v.path + " required by " + parent.name + " is missing initializer", null);
+	}
+	
+	function collect( cur : ShaderInfos, out : Array<ShaderInfos>, vertex : Bool ) {
+		if( cur.onStack )
+			error("Loop in shader dependencies ("+cur.name+")", null);
+		if( cur.marked == vertex )
+			return;
+		cur.marked = vertex;
+		cur.onStack = true;
+		for( d in cur.deps.keys() )
+			collect(d, out, vertex);
+		if( cur.vertex == vertex )
+			out.push(cur);
+		cur.onStack = false;
 	}
 	
-	public function link( shaders : Array<ShaderData>, startFun : String, outVar : String ) : ShaderData {
+	public function link( shadersData : Array<ShaderData>, outVars : Array<String> ) : ShaderData {
 		varMap = new Map();
 		allVars = new Array();
+		allFuns = new Map();
+		shaders = [];
+		
 		// globalize vars
-		for( s in shaders ) {
-			for( v in s.vars )
-				allocVar(v, null, null, null);
+		for( s in shadersData ) {
+			for( v in s.vars ) {
+				var kind = switch( v.name ) {
+				case "vertex": Vertex;
+				case "fragment": Fragment;
+				case "__init__": Init;
+				default: Function;
+				}
+				allocVar(v, null).kind = kind;
+			}
+		}
+		
+		// create shader segments
+		var priority = 0;
+		for( s in shadersData ) {
+			for( f in s.funs ) {
+				var v = allocVar(f.ref, f.expr.p);
+				switch( v.kind ) {
+				case Vertex, Fragment:
+					addShader(s.name+"."+(v.kind == Vertex ? "vertex" : "fragment"),v.kind == Vertex,f.expr, priority);
+				case Init:
+					switch( f.expr.e ) {
+					case TBlock(el):
+						var index = 0;
+						for( e in el )
+							addShader(s.name+".__init__"+(index++),true,e, -1);
+					default:
+						addShader(s.name+".__init__",true,f.expr, -1);
+					}
+				case Function:
+					allFuns.set(v.id, {
+						ref : v.v,
+						args : f.args,
+						ret : f.ret,
+						expr : mapExprVar(f.expr),
+					});
+				case Var:
+					throw "assert";
+				}
+			}
+			priority++;
+		}
+		shaders.sort(sortByPriorityDesc);
+		
+		// build dependency tree
+		var s = new ShaderInfos("", true);
+		s.deps = new Map();
+		for( outVar in outVars ) {
+			var v = varMap.get(outVar);
+			if( v == null )
+				throw "Variable not found " + outVar;
+			buildDependency(s, v);
+		}
+		
+		// collect needed dependencies
+		var v = [], f = [];
+		collect(s, v, true);
+		collect(s, f, false);
+		if( v.pop().name != "" ) throw "assert";
+		
+		// check that all dependencies are matched
+		for( s in shaders )
+			s.marked = null;
+		for( s in v.concat(f) ) {
+			for( d in s.deps.keys() )
+				if( d.marked == null )
+					error(d.name + " needed by " + s.name + " is unreachable", null);
+			s.marked = true;
+		}
+		
+		// build resulting vars
+		var outVars = [];
+		var varMap = new Map();
+		function addVar(v:AllocatedVar) {
+			if( varMap.exists(v.id) )
+				return;
+			varMap.set(v.id, true);
+			if( v.v.parent != null )
+				addVar(v.parent);
+			else
+				outVars.push(v.v);
+		}
+		for( s in v.concat(f) ) {
+			for( v in s.read )
+				addVar(v);
+			for( v in s.write )
+				addVar(v);
+		}
+		// build resulting shader functions
+		function build(name, a:Array<ShaderInfos> ) : TFunction {
+			var v : TVar = {
+				name : name,
+				type : TFun([ { ret : TVoid, args : [] } ]),
+				kind : Function,
+			};
+			outVars.push(v);
+			var exprs = [];
+			for( s in a )
+				switch( s.body.e ) {
+				case TBlock(el):
+					for( e in el ) exprs.push(e);
+				default:
+					exprs.push(s.body);
+				}
+			return {
+				ref : v,
+				ret : TVoid,
+				args : [],
+				expr : { e : TBlock(exprs), t : TVoid, p : exprs[0].p },
+			};
 		}
-		var vars = [];
-		for( v in allVars )
-			if( v.v.parent == null )
-				vars.push(v.v);
-		return { vars : vars, funs : [] };
+		var funs = [
+			build("vertex", v),
+			build("fragment", f),
+		];
+		return { name : "out", vars : outVars, funs : funs };
 	}
 	
 }

test/Test.hx

@@ -188,9 +188,7 @@ static var SRC = {
 	@param var power : Float;
 	@param var size : Float;
 	
-	var output : {
-		var color : Vec4;
-	};
+	var pixelColor : Vec4;
 	
 	@global var camera : {
 		var projDiag : Vec3;
@@ -205,7 +203,7 @@ static var SRC = {
 	
 	function fragment() {
 		var e = 1. - transformedNormal.normalize().dot(camera.dir.normalize());
-		output.color = color * e.pow(power);
+		pixelColor = color * e.pow(power);
 	}
 
 }}
@@ -332,9 +330,11 @@ class Test {
 	static function main() {
 		var shaders = [
 			new Proto().compile(),
-			new Outline().compile(),
+			new Texture().compile(),
+			new AnimatedUV().compile(),
+//			new Outline().compile(),
 		];
-		var s = new hxsl.Linker().link(shaders, "vertex", "out.color");
+		var s = new hxsl.Linker().link(shaders, ["output.position", "output.color"]);
 		trace("\n"+hxsl.Printer.shaderToString(s));
 	}