GLua code for performing mesh-based tracelines

supertrace.lua

-- "Supertrace" mesh tracing library 
-- (CC0 Public Domain, free to use in whatever you want)

-- Created by: Zachary Blystone ( [email protected] )

--[[
    This library performs a trace against an entity's visual meshes.
    While it is not as optimized as it could be, you won't see
    much of a framedrop on meshes with less than 5,000 triangles.

    Meshes are stored in a LRU cache to share loaded mesh data across
    all entities which have that model loaded. If you are using this
    on a lot of entities simultaneously with different models and you're
    seeing a lot of hitching, consider increasing the LRU_CACHE_SIZE.

    KNOWN BUGS:
    This sytem does not work with skinned meshes such as ragdolls.
    Some props won't work correctly because their mesh is offset via a
    root bone.

    USAGE:
    -- Create a supertrace object for a given entity
    local super_trace = supertrace.New( <entity> )

    -- Call super_trace:Trace with the ray position and forward vector
    local tr = super_trace:Trace( EyePos(), EyeAngles():Forward() )

    -- Calling super_trace:DrawHitTriangle will draw the triangle that was hit
    -- Calling super_trace:DrawMeshes will draw the meshes on the object

    -- tr contains a TraceResult-like structure:
    -- tr.Entity [will always be the entity this supertrace uses]
    -- tr.Hit [true if the trace hit]
    -- tr.Distance [distance of hitpos along ray]
    -- tr.HitPos [position that was hit]
    -- tr.HitNormal [interpolated normal of the hit triangle]
    -- tr.HitTexCoord [interpolated uv coordinates of triangle that was hit]
    -- tr.HitMesh [the index of the mesh that was hit]
    -- tr.HitTri [the index of the triangle that was hit (within mesh)]
    -- tr.HitMaterial = [IMaterial on the mesh that was hit]
]]

module("supertrace", package.seeall)

-- Set this to true if you'd like to see a demonstration
ENABLE_DEMONSTRATION = false

-- Change size of LRU cache
LRU_CACHE_SIZE = 16

-- Local functions to reduce indexing churn
local __vtab = FindMetaTable("Vector")
local __vset = __vtab.Set
local __vadd = __vtab.Add
local __vsub = __vtab.Sub
local __vdot = __vtab.Dot
local __vmul = __vtab.Mul
local __vlensqr = __vtab.LengthSqr
local __vdistsqr = __vtab.DistToSqr
local __vnorm = __vtab.Normalize
local __vunpack = __vtab.Unpack
local __vsetunpacked = __vtab.SetUnpacked

-- Transform the given normal into the coordinate space of the given matrix
local function TransformNormal( mtx, normal )

	local A00, A01, A02, A03,
		  A10, A11, A12, A13,
		  A20, A21, A22, A23,
		  A30, A31, A32, A33 = mtx:Unpack()

	local x,y,z = normal:Unpack()

	return Vector(
		A00 * x + A01 * y + A02 * z,
		A10 * x + A11 * y + A12 * z,
		A20 * x + A21 * y + A22 * z
	)

end

-- Returns the normal of the given triangle vertices
-- Note: this returns a shared vector
local __vtemp = Vector()
local __vtemp2 = Vector()
local __vtemp3 = Vector()
local function TriNormal(a,b,c)

	__vset(__vtemp, a)
	__vsub(__vtemp, b)
	__vset(__vtemp2, c)
	__vsub(__vtemp2, a)

    -- Cross-product creates a vector (more GC), do it manually here
    local a1,a2,a3 = __vunpack(__vtemp)
    local b1,b2,b3 = __vunpack(__vtemp2)

    __vsetunpacked(__vtemp3,
        a2 * b3 - a3 * b2,
        a3 * b1 - a1 * b3,
        a1 * b2 - a2 * b1
    )

	__vnorm(__vtemp3)
	return __vtemp3

end

-- Calculates the Barycentric coordinates of a point within a triangle
local __vedge1 = Vector()
local __vedge2 = Vector()
local __vedge3 = Vector()
local function TriBarycentric(a,b,c,point)

    __vset(__vedge1, b)
    __vsub(__vedge1, a)
    __vset(__vedge2, c)
    __vsub(__vedge2, a)
    __vset(__vedge3, point)
    __vsub(__vedge3, a)
	local v0 = __vedge1
	local v1 = __vedge2
	local v2 = __vedge3
	local d00 = __vdot(v0, v0)
	local d01 = __vdot(v0, v1)
	local d11 = __vdot(v1, v1)
	local d20 = __vdot(v2, v0)
	local d21 = __vdot(v2, v1)
	local denom = d00 * d11 - d01 * d01
	if denom == 0 then return 0,0,0 end

    local idenom = 1 / denom
	local v = (d11 * d20 - d01 * d21) * idenom
	local w = (d00 * d21 - d01 * d20) * idenom
	local u = 1.0 - v - w
	return u,v,w

end

-- Computes the time-of-intersection of a ray v plane
-- Not using util.IntersectRayWithPlane because it creates a new Vector (GC)
-- and this gets called a lot.
local __vraytemp = Vector()
local function RayIntersectPlane( normal, origin, p, dir )

	local d = __vdot(normal, dir)
	if d < -0.0001 then

        __vset(__vraytemp, origin)
        __vsub(__vraytemp, p)
		t = __vdot(__vraytemp, normal) / d
		return t >= 0, t

	end

	return false

end

-- Computes the intersection between a ray and a single triangle
-- Returns the hit position and its Barycentric coordinates
local __eptemp = Vector()
local function RayIntersectTriangle( tris, i, p, dir )

	local A = tris[i].pos
	local B = tris[i+1].pos
	local C = tris[i+2].pos

    -- Compute the triangle normal, early-out if facing wrong way
	local N = TriNormal(A,B,C)
	if __vdot(N, dir) > 0 then return false end

    -- Compute plane intersection against plane triangle sits on
	local check, t = RayIntersectPlane( N, A, p, dir )
	if not check then return false end

    -- ep = p + dir * t
    __vset(__eptemp, dir)
    __vmul(__eptemp, t)
    __vadd(__eptemp, p)
    local ep = __eptemp

    -- Compute Barycentric coordinates of hit point
    local u,v,w = TriBarycentric(A, B, C, ep)

    -- If Barycentric coordinates rest outside of triangle, it's not a hit
    if u == 0 and v == 0 and w == 0 then return false end
    if u < 0 or u > 1 or v < 0 or v > 1 or w < 0 or w > 1 then return false end

    return true, __eptemp, u,v,w

end

-- Computes the intersection between a ray and a soup of triangles
-- Returns the triangle index, hit position and its Barycentric coordinates
local function RayIntersectTriangles( tris, p, dir )

    local dist = math.huge
    local hit, pos, hit_i, u, v, w = false
	for i=#tris - 2, 1, -3 do

        -- Test each triangle (could optimize via a data structure [BVH / OT])
		local _hit, _pos, _u, _v, _w = RayIntersectTriangle( tris, i, p, dir )
		if _hit then
            hit = true
            local _dist = __vdistsqr(_pos, p) 
            if _dist < dist then
                -- Store if this triangle is closer than the last hit
                pos, hit_i, u, v, w = _pos, i, _u, _v, _w
                dist = _dist
            end
        end

	end

    -- If hit, copy the shard position vector
    if hit then pos = Vector(pos) end
	return hit, pos, hit_i, u, v, w 

end

-- Supertrace metatable
local meta = {}
meta.__index = meta

-- LRU cache for mesh data, older least-recently-used data will get overwritten
local mesh_lru_cache = {}
for i=1, LRU_CACHE_SIZE do mesh_lru_cache[#mesh_lru_cache+1] = {} end

function meta:Init( entity )

    assert( IsValid(entity), "SuperTrace: Entity expected" )

    -- Try to find cached mesh data
    local meshes = nil
    local model = entity:GetModel()
    for k,v in ipairs(mesh_lru_cache) do
        if v.model == model then
            meshes = v.meshes
            table.remove(mesh_lru_cache, k)
            mesh_lru_cache[#mesh_lru_cache+1] = v
            break
        end
    end

    -- Not cached, load the mesh data
    if meshes == nil then

        meshes = util.GetModelMeshes(model)
        for _, m in ipairs(meshes) do
            m.material = Material(m.material)
        end
        local cache = mesh_lru_cache[1]
        table.remove(mesh_lru_cache, 1)
        cache.meshes = meshes
        cache.model = model
        mesh_lru_cache[#mesh_lru_cache+1] = cache

    end

    assert(meshes ~= nil, "Unable to load meshes for : " .. tostring(model))

    self.entity = entity
    self.meshes = meshes
    self.result = {
        HitNormal = Vector(),
        HitTexCoord = {0,0},
    }
    self.mtx = Matrix()
    self.imtx = Matrix()
    self.old_pos = Vector()
    self.old_ang = Vector()
    self.matrices_valid = false
    return self

end

-- Update matrices from the entity if the entity has moved at all
function meta:UpdateMatrices()

    if not IsValid(self.entity) then return end

    local pos, ang = self.entity:GetPos(), self.entity:GetAngles()

    if pos ~= self.old_pos or ang ~= self.old_ang then

        self.mtx:SetTranslation( pos )
        self.mtx:SetAngles( ang )
        self.imtx:SetTranslation( pos )
        self.imtx:SetAngles( ang )
        self.matrices_valid = self.imtx:Invert()

        self.old_pos = pos
        self.old_ang = ang

    end

end

-- Draws the meshes on the entity
function meta:DrawMeshes()

    self:UpdateMatrices()

    cam.PushModelMatrix( self.mtx )

    for k, m in ipairs(self.meshes) do

        local c = Color(255,(k-1)*100,0,80)

        local tris = m.triangles
        for i=1, #tris, 3 do

            local v0 = tris[i]
            local v1 = tris[i+1]
            local v2 = tris[i+2]
    
            render.DrawLine(v0.pos, v1.pos, c)
            render.DrawLine(v1.pos, v2.pos, c)
            render.DrawLine(v2.pos, v0.pos, c)
    
        end
    end

    cam.PopModelMatrix()

end

-- Draws the triangle that was hit
function meta:DrawHitTriangle()

    local res = self.result
    if not res.Hit then return end

    cam.PushModelMatrix( self.mtx )

    local tris = self.meshes[res.HitMesh].triangles
    local i = res.HitTri

    local v0 = tris[i]
    local v1 = tris[i+1]
    local v2 = tris[i+2]

    render.DrawLine(v0.pos, v1.pos, c)
    render.DrawLine(v1.pos, v2.pos, c)
    render.DrawLine(v2.pos, v0.pos, c)

    cam.PopModelMatrix()

end

-- Sums the products of each vector with its equivelent barycentric component
local __vmultemp0 = Vector()
local __vmultemp1 = Vector()
local __vmultemp2 = Vector()
local __vmulsum = Vector()
local function BarycentricMul(v0, v1, v2, u, v, w)

    __vset(__vmultemp0, v0)
    __vmul(__vmultemp0, u)
    __vset(__vmultemp1, v1)
    __vmul(__vmultemp1, v)
    __vset(__vmultemp2, v2)
    __vmul(__vmultemp2, w)
    __vset(__vmulsum, __vmultemp0)
    __vadd(__vmulsum, __vmultemp1)
    __vadd(__vmulsum, __vmultemp2)
    return __vmulsum

end

local __vdist = Vector()
function meta:Trace( origin, dir )

    self:UpdateMatrices()

    self.result.Hit = false

    if not self.matrices_valid then return self.result end

    local res = self.result

    origin = self.imtx * origin
	dir = TransformNormal( self.imtx, dir )

    local best_distance = math.huge
    for k, m in ipairs(self.meshes) do

        local tris = m.triangles
        local hit, pos, i, u, v, w = RayIntersectTriangles(tris, origin, dir)
        if hit then
            __vset(__vdist, pos)
            __vsub(__vdist, origin)
            local distance = __vdot(__vdist, dir)
            if distance <= best_distance then
                local v0 = tris[i]
                local v1 = tris[i+1]
                local v2 = tris[i+2]
                local vnormal = BarycentricMul(
                    v0.normal, 
                    v1.normal, 
                    v2.normal,
                    u, v, w)

                res.Entity = self.entity
                res.Hit = true
                res.Distance = distance
                res.HitPos = pos
                res.HitNormal = vnormal
                res.HitTexCoord[1] = v0.u * u + v1.u * v + v2.u * w
                res.HitTexCoord[2] = v0.v * u + v1.v * v + v2.v * w
                res.HitMesh = k
                res.HitTri = i
                res.HitMaterial = m.material
                best_distance = distance
            end

        end

    end

    if res.Hit then
        res.HitPos = self.mtx * res.HitPos
        res.HitNormal = TransformNormal( self.mtx, res.HitNormal )
    end

    return res

end

function New( entity )
    return setmetatable({}, meta):Init(entity)
end

if ENABLE_DEMONSTRATION then

    local uv_material = CreateMaterial("uv_mat2", "UnLitGeneric", {
        ["$vertexcolor"] = "1",
        ["$vertexalpha"] = "0",
    })

    local function Render()

        -- Normal trace
        local tr = LocalPlayer():GetEyeTrace()
        if not tr.Hit or not IsValid(tr.Entity) then return end

        -- Super trace on hit entity
        local super_trace = supertrace.New(tr.Entity)

        -- Perform supertrace
        local tr = super_trace:Trace( EyePos(), EyeAngles():Forward() )

        -- Exit if trace doesn't hit
        if not tr.Hit then return end
        
        -- Draw the hit triangle
        super_trace:DrawHitTriangle()

        -- Draw the UV coordinate
        local mat = tr.HitMaterial
        local tex = mat:GetTexture("$basetexture")

        uv_material:SetTexture("$basetexture", tex)

        cam.Start2D()

        surface.SetMaterial(uv_material)
        surface.SetDrawColor(200,200,200,255)
        surface.DrawTexturedRect(0,0, tex:Width(),tex:Height())

        surface.SetDrawColor(0,0,0,255)
        surface.DrawRect(
            tr.HitTexCoord[1] * tex:Width() - 4, 
            tr.HitTexCoord[2] * tex:Height() - 4, 
            8, 8)

        surface.SetDrawColor(255,255,255,255)
        surface.DrawRect(
            tr.HitTexCoord[1] * tex:Width() - 2, 
            tr.HitTexCoord[2] * tex:Height() - 2, 
            4, 4)

        cam.End2D()

    end

    hook.Add("PostDrawOpaqueRenderables", "supertrace", Render)

else

    hook.Remove("PostDrawOpaqueRenderables", "supertrace")

end