Grass: add crude instancing code

1 files changed, 343 insertions, 0 deletions

diff --git a/Scripts/InstanceGrass.cs b/Scripts/InstanceGrass.cs
new file mode 100755
index 0000000..3c0d8f6
--- /dev/null
+++ b/Scripts/InstanceGrass.cs
@@ -0,0 +1,343 @@
+using UdonSharp;
+using UnityEngine;
+using VRC.SDKBase;
+
+#if UNITY_EDITOR && !COMPILER_UDONSHARP
+using UnityEditor;
+using System.IO;
+#endif
+
+public class InstanceGrass : UdonSharpBehaviour
+{
+  [SerializeField] public GameObject prefab_;
+  // The density of instances, in instances per meter.
+  [SerializeField] public float density_;
+  // The extent along each cardinal axis where instances will be rendered, in
+  // meters. I.e. render inside a cube with edges this long.
+  [SerializeField] public Vector3 extent_meters_;
+  [SerializeField] public Vector3 angle_randomization_;
+  [SerializeField] public float scale_randomization_;
+
+  // GPU mode settings
+  [SerializeField] public Vector3Int max_grid_size_ = new Vector3Int(512, 1, 512);  // Maximum grid dimensions for GPU mode
+  [SerializeField] public Material blit_material_;  // Auto-generated in editor, set by custom editor
+
+  private Mesh mesh_;
+  private Material instance_material_;  // Extracted from prefab
+  private Transform base_transform_;
+  private Vector3 cell_dim_;
+  private Vector3 inv_cell_dim_;
+  private Vector3Int count_;
+  private int total_instances_;
+  private bool initialized_;
+
+  // Track fields to detect runtime changes.
+  private float density_live_;
+  private Vector3 extent_meters_live_;
+  private Vector3 angle_randomization_live_;
+  private float scale_randomization_live_;
+
+  // GPU-specific resources
+  private RenderTexture instance_data_tex_;
+  private Matrix4x4[] identity_transforms_;  // All identity matrices
+  private int tex_width_;
+  private int tex_height_;
+
+  private void Init() {
+    // Extract components from prefab.
+    if (prefab_ != null) {
+      MeshFilter mesh_filter = prefab_.GetComponent<MeshFilter>();
+      if (mesh_filter == null && prefab_.transform.childCount > 0) {
+        mesh_filter = prefab_.GetComponentInChildren<MeshFilter>();
+      }
+      if (mesh_filter != null) {
+        mesh_ = mesh_filter.sharedMesh;
+        if (mesh_ == null) {
+          mesh_ = mesh_filter.mesh;
+        }
+      } else {
+        Debug.LogError("[Grass::Debug] Could not find MeshFilter on prefab or children.");
+      }
+
+      MeshRenderer mesh_renderer = prefab_.GetComponent<MeshRenderer>();
+      if (mesh_renderer == null && prefab_.transform.childCount > 0) {
+        mesh_renderer = prefab_.GetComponentInChildren<MeshRenderer>();
+      }
+      if (mesh_renderer != null) {
+        Material[] materials = mesh_renderer.sharedMaterials;
+        if (materials != null && materials.Length > 0) {
+          instance_material_ = materials[0];
+        } else {
+          Debug.LogError("[Grass::Debug] MeshRenderer has no materials.");
+        }
+      } else {
+        Debug.LogError("[Grass::Debug] Could not find MeshRenderer on prefab.");
+      }
+
+      base_transform_ = prefab_.transform;
+      Debug.Log($"[Grass::Debug] Prefab transform - pos:{base_transform_.position}, rot:{base_transform_.rotation.eulerAngles}, scale:{base_transform_.localScale}");
+    } else {
+      Debug.LogError("[Grass::Debug] prefab is null in Init().");
+    }
+
+    density_ = Mathf.Max(1e-6f, density_);
+    extent_meters_ = Vector3.Max(Vector3.one * 1e-6f, extent_meters_);
+    angle_randomization_ = new Vector3(
+        Mathf.Clamp(angle_randomization_.x, 0f, 180f),
+        Mathf.Clamp(angle_randomization_.y, 0f, 180f),
+        Mathf.Clamp(angle_randomization_.z, 0f, 180f));
+
+    // Use max_grid_size_ directly for GPU mode
+    count_ = Vector3Int.Min(max_grid_size_, new Vector3Int(
+        Mathf.Max(1, Mathf.RoundToInt(density_ * extent_meters_.x)),
+        Mathf.Max(1, Mathf.RoundToInt(density_ * extent_meters_.y)),
+        Mathf.Max(1, Mathf.RoundToInt(density_ * extent_meters_.z))));
+
+    cell_dim_ = new Vector3(
+        extent_meters_.x / count_.x,
+        extent_meters_.y / count_.y,
+        extent_meters_.z / count_.z);
+    inv_cell_dim_ = new Vector3(
+        1f / cell_dim_.x,
+        1f / cell_dim_.y,
+        1f / cell_dim_.z);
+
+    total_instances_ = count_.x * count_.y * count_.z;
+
+    // Calculate texture dimensions (power of 2, minimum 256x256)
+    int min_tex_size = Mathf.CeilToInt(Mathf.Sqrt(total_instances_));
+    tex_width_ = Mathf.Max(256, Mathf.NextPowerOfTwo(min_tex_size));
+    tex_height_ = tex_width_;
+
+    // Create instance data texture
+    if (instance_data_tex_ != null) {
+      instance_data_tex_.Release();
+    }
+    instance_data_tex_ = new RenderTexture(tex_width_, tex_height_, 0, RenderTextureFormat.ARGBFloat);
+    instance_data_tex_.filterMode = FilterMode.Point;
+    instance_data_tex_.Create();
+
+    // Create transform array with encoded instance IDs - must match actual draw count
+    int actual_draw_count = tex_width_ * tex_height_;
+    if (identity_transforms_ == null || identity_transforms_.Length != actual_draw_count) {
+      identity_transforms_ = new Matrix4x4[actual_draw_count];
+      for (int i = 0; i < actual_draw_count; i++) {
+        // Encode instance ID in the matrix translation component
+        // We'll put the ID in the X component, which the shader will read
+        Matrix4x4 m = Matrix4x4.identity;
+        m.m03 = i; // Store instance ID in translation.x
+        identity_transforms_[i] = m;
+      }
+    }
+
+    if (mesh_ != null) {
+      // Set huge bounds so culling doesn't interfere
+      mesh_.bounds = new Bounds(Vector3.zero, Vector3.one * 10000f);
+    }
+
+    density_live_ = density_;
+    extent_meters_live_ = extent_meters_;
+    angle_randomization_live_ = angle_randomization_;
+    scale_randomization_live_ = scale_randomization_;
+
+    Debug.Log($"[Grass::Debug] Init: density={density_}, extent={extent_meters_}, count={count_}, cell_dim={cell_dim_}, instances={total_instances_}, tex={tex_width_}x{tex_height_}, scale={base_transform_.localScale}, rot={base_transform_.localRotation.eulerAngles}");
+    Debug.Log($"[Grass::Debug] Init details: count_.x={count_.x}, count_.y={count_.y}, count_.z={count_.z}, extent.x={extent_meters_.x}, extent.z={extent_meters_.z}");
+  }
+
+  private bool Valid() {
+    if (prefab_ == null) {
+      Debug.LogError("[Grass::Debug] prefab is null.");
+      return false;
+    }
+    if (mesh_ == null) {
+      Debug.LogError("[Grass::Debug] mesh is null.");
+      return false;
+    }
+    if (instance_material_ == null) {
+      Debug.LogError("[Grass::Debug] instance_material is null.");
+      return false;
+    }
+    if (blit_material_ == null) {
+      Debug.LogError("[Grass::Debug] blit_material is null (failed to generate).");
+      return false;
+    }
+    if (instance_data_tex_ == null) {
+      Debug.LogError("[Grass::Debug] instance_data_tex is null.");
+      return false;
+    }
+    return true;
+  }
+
+  void Start() {
+    Init();
+    if (!Valid()) {
+      return;
+    }
+
+    initialized_ = true;
+    Debug.Log($"[Grass::Debug] GPU mode initialized");
+  }
+
+  void Update() {
+    // Reinitialize if any config changed at runtime.
+    if (density_ != density_live_ ||
+        extent_meters_ != extent_meters_live_ ||
+        angle_randomization_ != angle_randomization_live_ ||
+        scale_randomization_ != scale_randomization_live_) {
+      Init();
+      initialized_ = false;
+    }
+
+    if (!Valid()) {
+      return;
+    }
+
+    VRCPlayerApi lcl_player = Networking.LocalPlayer;
+    Vector3 player_pos = lcl_player.GetPosition();
+    int grid_x = Mathf.FloorToInt(player_pos.x * inv_cell_dim_.x);
+    int grid_y = Mathf.FloorToInt(player_pos.y * inv_cell_dim_.y);
+    int grid_z = Mathf.FloorToInt(player_pos.z * inv_cell_dim_.z);
+
+    int half_x = count_.x / 2;
+    int half_y = count_.y / 2;
+    int half_z = count_.z / 2;
+
+    // Update blit material properties
+    blit_material_.SetVector("_PlayerGridPos", new Vector3(grid_x, grid_y, grid_z));
+    blit_material_.SetVector("_GridCount", new Vector3(count_.x, count_.y, count_.z));
+    blit_material_.SetVector("_GridHalf", new Vector3(half_x, half_y, half_z));
+    blit_material_.SetVector("_TexDimensions", new Vector2(tex_width_, tex_height_));
+
+    // Blit to generate instance data texture
+    VRCGraphics.Blit(null, instance_data_tex_, blit_material_);
+
+    // Update instance material properties
+    instance_material_.SetTexture("_Instance_Texture_Offset_Data_Tex", instance_data_tex_);
+    instance_material_.SetVector("_Instance_Texture_Offset_Cell_Dimensions", cell_dim_);
+    instance_material_.SetVector("_Instance_Texture_Offset_Angle_Randomization", angle_randomization_);
+    instance_material_.SetFloat("_Instance_Texture_Offset_Scale_Randomization", scale_randomization_);
+    instance_material_.SetVector("_Instance_Texture_Offset_Base_Scale", base_transform_.localScale);
+
+    // Pass rotation as quaternion (x, y, z, w)
+    Quaternion rot = base_transform_.localRotation;
+    instance_material_.SetVector("_Instance_Texture_Offset_Base_Rotation", new Vector4(rot.x, rot.y, rot.z, rot.w));
+
+    if (Time.frameCount % 300 == 0) {
+      Debug.Log($"[Grass::Debug] Drawing {total_instances_} GPU instances");
+      Debug.Log($"[Grass::Debug] GridCount={count_}, TexDim={tex_width_}x{tex_height_}, CellDim={cell_dim_}");
+    }
+
+    // Draw instances with identity transforms - GPU handles everything
+    // Must draw tex_width * tex_height instances to match texture
+    VRCGraphics.DrawMeshInstanced(
+        mesh_,
+        0,
+        instance_material_,
+        identity_transforms_,
+        total_instances_,
+        null,
+        UnityEngine.Rendering.ShadowCastingMode.Off,
+        true,
+        0,
+        null,
+        UnityEngine.Rendering.LightProbeUsage.Off,
+        null);
+
+    initialized_ = true;
+  }
+
+  void OnDestroy() {
+    if (instance_data_tex_ != null) {
+      instance_data_tex_.Release();
+    }
+  }
+}
+
+#if UNITY_EDITOR && !COMPILER_UDONSHARP
+[CustomEditor(typeof(InstanceGrass))]
+public class InstanceGrassEditor : Editor
+{
+  public override void OnInspectorGUI()
+  {
+    DrawDefaultInspector();
+
+    InstanceGrass script = (InstanceGrass)target;
+
+    // Auto-generate blit material if missing
+    if (script.blit_material_ == null)
+    {
+      if (GUILayout.Button("Generate Blit Material"))
+      {
+        GenerateBlitMaterial(script);
+      }
+
+      EditorGUILayout.HelpBox("Blit material is missing. Click the button above to auto-generate it.", MessageType.Warning);
+    }
+    else
+    {
+      EditorGUILayout.HelpBox("Blit material is set. GPU instancing ready.", MessageType.Info);
+    }
+  }
+
+  private void GenerateBlitMaterial(InstanceGrass script)
+  {
+    Shader blitShader = Shader.Find("yum_food/GrassGridBlit");
+    if (blitShader == null)
+    {
+      EditorUtility.DisplayDialog("Error", "Could not find shader 'yum_food/GrassGridBlit'. Make sure GrassGridBlit.shader is imported.", "OK");
+      return;
+    }
+
+    string assetPath = "Assets/yum_food/3ner/Grass_Generated/Grass_generated.mat";
+
+    // Create or load existing material
+    Material blitMat = AssetDatabase.LoadAssetAtPath<Material>(assetPath);
+    if (blitMat == null)
+    {
+      // Create new material
+      blitMat = new Material(blitShader);
+      blitMat.name = "Grass_generated";
+
+      // Ensure directory exists
+      string directory = Path.GetDirectoryName(assetPath);
+      if (!Directory.Exists(directory))
+      {
+        Directory.CreateDirectory(directory);
+      }
+
+      AssetDatabase.CreateAsset(blitMat, assetPath);
+      Debug.Log($"[Grass::Editor] Created blit material at {assetPath}");
+    }
+    else
+    {
+      // Update existing material shader
+      blitMat.shader = blitShader;
+      EditorUtility.SetDirty(blitMat);
+      Debug.Log($"[Grass::Editor] Updated existing blit material at {assetPath}");
+    }
+
+    script.blit_material_ = blitMat;
+    EditorUtility.SetDirty(script);
+    AssetDatabase.SaveAssets();
+  }
+
+  [UnityEditor.Callbacks.DidReloadScripts]
+  private static void OnScriptsReloaded()
+  {
+    // Auto-generate blit material for all InstanceGrass components when scripts reload
+    InstanceGrass[] allGrass = FindObjectsOfType<InstanceGrass>();
+    foreach (var grass in allGrass)
+    {
+      if (grass.blit_material_ == null)
+      {
+        var editor = CreateEditor(grass) as InstanceGrassEditor;
+        if (editor != null)
+        {
+          editor.GenerateBlitMaterial(grass);
+          DestroyImmediate(editor);
+        }
+      }
+    }
+  }
+}
+#endif