Sprite Batching with stage3D

SpriteBatch3D.as

package flash.display3D
{
	import com.adobe.utils.AGALMiniAssembler;
	
	import flash.display.Sprite;
	import flash.display3D.textures.*;
	import flash.errors.*;
	import flash.geom.*;

	/**
	 * 
	 *
	 * 
	 */	
	public class SpriteBatch3D
	{
		private var _context3D:Context3D;
		private var _matrix3D:Matrix3D;
		private var _program3D:Program3D;
		
		private var _vertexBuffer:VertexBuffer3D;
		private var _indexBuffer:IndexBuffer3D;
		
		private var _sprites:Vector.<SpriteBatch3DSprite>;
		private var _lock:Boolean;
		
		private static const DATA_PER_VERTEX:int = 4;
		private static const BATCH_SIZE:int = 1024;
		// TODO; Batch size can be adaptive instead
		
		public function SpriteBatch3D(context3D:Context3D) {
			if (context3D == null) {
				throw new ArgumentError("context3D cannot be null.");
			}
			
			_context3D = context3D;
			_program3D = createDefaultProgram();
			_sprites = new Vector.<SpriteBatch3DSprite>;
			
			var vertexData:Vector.<Number> = new Vector.<Number>(BATCH_SIZE * 4 * DATA_PER_VERTEX);
			_vertexBuffer = _context3D.createVertexBuffer(vertexData.length / DATA_PER_VERTEX, DATA_PER_VERTEX);
			_vertexBuffer.uploadFromVector(vertexData, 0, vertexData.length / DATA_PER_VERTEX);
			
			var indexData:Vector.<uint>    = new Vector.<uint>;
			for (var i:int = 0; i < BATCH_SIZE; i++) 
			{
				var n:int = i * 4;
				indexData.push(n, n + 1, n + 2,
					n + 0, n + 3, n + 2);
			}
			_indexBuffer = _context3D.createIndexBuffer(indexData.length);
			_indexBuffer.uploadFromVector(indexData, 0, indexData.length);
		}
		
		public function begin(matrix3D:Matrix3D=null):void {
			if (_lock) {
				throw new IllegalOperationError("already locked");
			}
			if(matrix3D == null) {
				matrix3D = new Matrix3D();
			}
			
			_matrix3D = matrix3D;
			_lock = true;
		}
		
		public function end():void {
			if (!_lock) {
				throw new IllegalOperationError("already unlocked"); 
			}
			
			if (_sprites.length < 0) {
				return;
			}
			
			_context3D.setProgram(_program3D);
			_context3D.setProgramConstantsFromMatrix(Context3DProgramType.VERTEX, 0, _matrix3D, true);
			_context3D.setVertexBufferAt(0, _vertexBuffer, 0, Context3DVertexBufferFormat.FLOAT_2); //x, y
			_context3D.setVertexBufferAt(1, _vertexBuffer, 2, Context3DVertexBufferFormat.FLOAT_2); //u, v
			
			var texture:Texture = _sprites[0].texture;
			var count:int = 0;
			
			for each(var sprite:SpriteBatch3DSprite in _sprites) {
					if (texture != sprite.texture || count == BATCH_SIZE) {
					_context3D.setTextureAt(0, texture);
					_context3D.drawTriangles(_indexBuffer, 0, count * 2);
					
					texture = sprite.texture;
					count = 0;
				}
					
				_vertexBuffer.uploadFromVector(sprite.vertices, count * 4, sprite.vertices.length / DATA_PER_VERTEX);
				
				count++;
			}
			
			if (count > 0) { 
				_context3D.setTextureAt(0, texture);
				_context3D.drawTriangles(_indexBuffer, 0, count * 2);
			}
			
			_context3D.setVertexBufferAt(0, null);
			_context3D.setVertexBufferAt(1, null);
			
			
			_sprites.splice(0, _sprites.length);
			_lock = false;
		}
		
		public function draw(texture:Texture, dst:Rectangle, src:Rectangle=null):void {
			if (!_lock) {
				throw new IllegalOperationError("no lock."); 
			}
			
			if (src == null) {
				src = new Rectangle(0, 0, 1, 1);
			}
			
			// TODO
			// add rotation, depth
			
			var sprite:SpriteBatch3DSprite = new SpriteBatch3DSprite(texture, Vector.<Number>([
				dst.left, 	dst.top, 	src.left, 	src.top,
				dst.left, 	dst.bottom, src.left, 	src.bottom,
				dst.right, 	dst.bottom, src.right, 	src.bottom,
				dst.right, 	dst.top, 	src.right, 	src.top]));
			
			_sprites.push(sprite);
		}

		private function createDefaultProgram():Program3D {	
			var vertexAssembler:AGALMiniAssembler = new AGALMiniAssembler();
			vertexAssembler.assemble(flash.display3D.Context3DProgramType.VERTEX,
				"m44 op, va0, vc0   \n" + // vertex * clipspace
				"mov v0, va1		\n"   // move uv
			);
			
			var fragmentAssembler:AGALMiniAssembler = new AGALMiniAssembler();
			fragmentAssembler.assemble(flash.display3D.Context3DProgramType.FRAGMENT, 
				"tex ft1, v0, fs0 <2d,linear,nomip> \n" +
				"mov oc, ft1"
			);
			
			var program3D:Program3D = _context3D.createProgram();
			program3D.upload(vertexAssembler.agalcode, fragmentAssembler.agalcode);
			
			return program3D;
		}
	}
}

SpriteBatch3DSprite.as

package flash.display3D
{
	import flash.display3D.textures.*;
	
	internal class SpriteBatch3DSprite {
		private var _texture:Texture;
		private var _vertices:Vector.<Number>;
		
		public function SpriteBatch3DSprite(texture:Texture, vertices:Vector.<Number>) {
			_texture  = texture;
			_vertices = vertices;
		}

		public function get texture():Texture
		{
			return _texture;
		}

		public function get vertices():Vector.<Number>
		{
			return _vertices;
		}
	}
}