BlenderSharp/BlendFile/Reader.cs

using System.Reflection;
using System.Text;
using Kaitai;


namespace BlendFile;

public class Reader {
    string _path;
    private readonly Dictionary<int, Type> dnaTypes = new();
    
    private Dictionary<(long, Type), object> objects = new();
    public Dictionary<(long, Type), object> Objects => objects;
    
    public List<object> GetObjects() => objects.Values.ToList();
    public List<T> GetObjects<T>() => objects.Values.OfType<T>().ToList();
    
    private SortedDictionary<long, Kaitai.BlendFile.FileBlock> memBlocks = new();
    
    /// <summary>
    /// Gets the block at the specified memory address
    /// </summary>
    /// <param name="memAddr">memory address in current system endianness</param>
    /// <returns>A <see cref="Kaitai.BlendFile.FileBlock"/> object</returns>
    public Kaitai.BlendFile.FileBlock? GetBlock(long memAddr) => memBlocks.SkipWhile(x => x.Key < memAddr).FirstOrDefault().Value;
    
    /// <summary>
    /// Creates a new instance of the <see cref="Reader"/> class
    /// </summary>
    /// <param name="path">A <see cref="string"/> containing a path to a blend file that will be read.</param>
    public Reader(string path) : this() {
        _path = path;
    }

    /// <summary>
    /// Creates a new instance of the <see cref="Reader"/> class
    /// </summary>
    public Reader() {
        _path = "";
        var types = Assembly.GetExecutingAssembly().DefinedTypes;
        foreach (var type in types) {
            var attrib = type.GetCustomAttribute<DNAClassAttribute>();
            if (attrib == null) continue;
            dnaTypes.Add(attrib.OriginalIndex, type);
        }
    }
    
    public void Read(string path) {
        _path = path;
        Read();
    }
    
    public void Read() {
        var file = new KaitaiStream(_path);
        var blend = new Kaitai.BlendFile(file);
        
        bool isLe = blend.Hdr.Endian == Kaitai.BlendFile.Endian.Le;
        //TODO: two blocks somehow have the same mem address... this sounds wrong.
        blend.Blocks.ForEach(block => memBlocks.TryAdd(block.MemAddr.ToMemAddr(isLe), block));
        
        foreach (var block in blend.Blocks)
        {
            //We need to read all blocks of data regardeless of the type
            //if (block.Code != "DATA") continue;
            //Checks if the block has a known SDNA type, meaning that we know how it is structured 
            if(!dnaTypes.ContainsKey((int)block.SdnaIndex)) continue;
            
            Type t = dnaTypes.Values.First(
                x => x.GetCustomAttribute<DNAClassAttribute>()!.OriginalName == block.SdnaStruct.Type);
            // How many objects are in the block
            var count = block.Count;
            //offset for the next object in the block
            var blockOffset = 0;
            //for each expected object in the block
            for(var i=0; i<count; i++) {
                //create an instance of type "t" (the dna type we inferred before)  
                var obj = Activator.CreateInstance(t);
                if(obj == null) continue; //should never happen
                
                //fill the object with the data from the block
                FillObject(block, ref obj, t.GetFields(), blockOffset);
                
                //move the offset to the next object in the block
                blockOffset += t.GetCustomAttribute<DNAClassAttribute>()!.Size;
            }
        }
    }

    /// <summary>
    /// 
    /// </summary>
    /// <param name="block"></param>
    /// <param name="obj"></param>
    /// <param name="fieldMetadata"></param>
    /// <param name="startOffset"></param>
    /// <returns></returns>
    private void FillObject(Kaitai.BlendFile.FileBlock block, ref object? obj, FieldInfo[] fieldMetadata, int startOffset = 0) {
        if(block.Code == "ENDB") return;
        foreach (var field in fieldMetadata) {
            //Get the DNAFieldAttribute of the current field
            var attrib = field.GetCustomAttribute<DNAFieldAttribute>();
            if (attrib == null) continue; //should never happen, but means a field has no metadata
            
            //Calculate the offset from where the data of the field starts.
            //Because the order of the fields is not guaranteed we need to compute it each time
            int offset = fieldMetadata.Where(f => f.GetCustomAttribute<DNAFieldAttribute>()!.OriginalIndex < attrib.OriginalIndex)
                .Sum(f => f.GetCustomAttribute<DNAFieldAttribute>()!.Size) + startOffset;
            
            int size = attrib.Size;
            var data = new byte[size];
            Array.Copy((byte[])block.Body, offset, data, 0, size);
            //Convert the data to the correct type
            object? value = ConvertFieldData(data, attrib.OriginalType);
            
            if(value == null){ //if the data could not be converted
                //Check if the field is a pointer to another DNA structure
                if (dnaTypes.Values.FirstOrDefault(x => x.GetCustomAttribute<DNAClassAttribute>()!.OriginalName == attrib.OriginalType) != null) {
                    //Create a new instance of the DNA structure type
                    object? newObj = Activator.CreateInstance(dnaTypes.Values.First(x =>
                        x.GetCustomAttribute<DNAClassAttribute>()!.OriginalName == attrib.OriginalType));
                    if(newObj == null) continue; //should never happen... type is missing?
                    
                    //Get the information of the fields of the new object
                    var fieldInfo = newObj.GetType().GetFields();
                    
                    //If the field is a pointer, we need to dereference it
                    if (attrib.IsPointer) {
                        //Get the memory address of the data
                        long addr = data.ToMemAddr();
                        //Get the closest block that contains the data
                        var derefBlock = GetBlock(addr);
                        //Calculate the relative address of the data in the block
                        long relAddr = addr - derefBlock!.MemAddr.ToMemAddr();
                        
                        if(objects.TryGetValue((addr, newObj.GetType()), out object? o)){ //If the object is already created, we can just assign it
                            field.SetValue(obj, o);
                            continue;
                        }
                        
                        //Fill the object with the data from the block (this is recursive)
                        FillObject(derefBlock, ref newObj, fieldInfo, (int)relAddr);
                    } else {
                        long relAddr = block.MemAddr.ToMemAddr() + offset;
                        if(objects.TryGetValue((relAddr, newObj.GetType()), out object? o)){ //If the object is already created, we can just assign it
                            field.SetValue(obj, o);
                            continue;
                        }
                        //Fill the object with the data from the block (this is recursive)
                        FillObject(block, ref newObj, fieldInfo, offset);
                    }
                }
                continue; //should never happen, but means the data could not be converted
            }
            //Additionally... some fields might not be nullable so it's better to not assign the value and leave the default one.
            field.SetValue(obj, value);
        }
        objects.Add((block.MemAddr.ToMemAddr() + startOffset, obj!.GetType()), obj!);
    }
    
    private object? ConvertFieldData(byte[] data, string type)
    {
        return type switch
        {
            "char" => (char)data[0],
            "short" => BitConverter.ToInt16(data, 0),
            "int" => BitConverter.ToInt32(data, 0),
            "float" => BitConverter.ToSingle(data, 0),
            "double" => BitConverter.ToDouble(data, 0),
            "string" => Encoding.UTF8.GetString(data), // utf8?
            "void" => null, // object?
            "ushort" => BitConverter.ToUInt16(data, 0),
            "uchar" => data[0],
            "int64_t" => BitConverter.ToInt64(data, 0),
            "int8_t" => (sbyte)data[0],
            "uint64_t" => BitConverter.ToUInt64(data, 0),
            _ => null
        };
    }
}