diff --git a/src/commonMain/kotlin/space/kscience/gdml/Gdml.kt b/src/commonMain/kotlin/space/kscience/gdml/Gdml.kt
index dc28a97..6d1cc62 100644
--- a/src/commonMain/kotlin/space/kscience/gdml/Gdml.kt
+++ b/src/commonMain/kotlin/space/kscience/gdml/Gdml.kt
@@ -111,7 +111,7 @@ public class Gdml : GdmlRegistry {
     public companion object
 }
 
-public inline fun Gdml(block: Gdml.() -> Unit): Gdml  = Gdml().apply(block)
+public inline fun Gdml(block: Gdml.() -> Unit): Gdml = Gdml().apply(block)
 
 public inline fun <reified T : GdmlDefine> GdmlRef<T>.resolve(root: Gdml): T? = root.getDefine(ref)
 public inline fun <reified T : GdmlSolid> GdmlRef<T>.resolve(root: Gdml): T? = root.getSolid(ref)
diff --git a/src/commonMain/kotlin/space/kscience/gdml/GdmlMaterialPostProcessor.kt b/src/commonMain/kotlin/space/kscience/gdml/GdmlMaterialPostProcessor.kt
index 4517a05..a7a784d 100644
--- a/src/commonMain/kotlin/space/kscience/gdml/GdmlMaterialPostProcessor.kt
+++ b/src/commonMain/kotlin/space/kscience/gdml/GdmlMaterialPostProcessor.kt
@@ -1,72 +1,68 @@
 package space.kscience.gdml
 
-public class GdmlMaterialPostProcessor {
-    public companion object {
-        public fun removeUnusedMaterials(gdml: Gdml): Gdml {
+public fun Gdml.removeUnusedMaterials(): Gdml {
 
-            val materialsInUse: MutableSet<GdmlMaterial> = mutableSetOf()
-            val elementsInUse: MutableSet<GdmlElement> = mutableSetOf()
-            val isotopesInUse: MutableSet<GdmlIsotope> = mutableSetOf()
+    val materialsInUse: MutableSet<GdmlMaterial> = mutableSetOf()
+    val elementsInUse: MutableSet<GdmlElement> = mutableSetOf()
+    val isotopesInUse: MutableSet<GdmlIsotope> = mutableSetOf()
 
-            fun addMaterialsRecursively(composite: GdmlComposite?) {
-                // composite may be composed of other composites
-                composite?.fractions?.forEach { fraction ->
-                    val materialComponent = gdml.materials.get<GdmlComposite>(fraction.ref)
-                    if (materialComponent != null) {
-                        materialsInUse.add(materialComponent)
-                        addMaterialsRecursively(materialComponent)
-                    } else {
-                        val elementComponent = gdml.materials.get<GdmlElement>(fraction.ref)
-                        if (elementComponent != null) {
-                            elementsInUse.add(elementComponent)
-                            elementComponent.fractions.forEach { elementFraction ->
-                                val isotopeComponent = gdml.materials.get<GdmlIsotope>(elementFraction.ref)
-                                isotopesInUse.add(isotopeComponent!!) // element must be mode of isotopes
-                            }
-                        } else {
-                            val isotopeComponent = gdml.materials.get<GdmlIsotope>(fraction.ref)
-                            if (isotopeComponent != null) {
-                                isotopesInUse.add(isotopeComponent)
-                            }
-                        }
+    fun addMaterialsRecursively(composite: GdmlComposite?) {
+        // composite may be composed of other composites
+        composite?.fractions?.forEach { fraction ->
+            val materialComponent = this.materials.get<GdmlComposite>(fraction.ref)
+            if (materialComponent != null) {
+                materialsInUse.add(materialComponent)
+                addMaterialsRecursively(materialComponent)
+            } else {
+                val elementComponent = this.materials.get<GdmlElement>(fraction.ref)
+                if (elementComponent != null) {
+                    elementsInUse.add(elementComponent)
+                    elementComponent.fractions.forEach { elementFraction ->
+                        val isotopeComponent = this.materials.get<GdmlIsotope>(elementFraction.ref)
+                        isotopesInUse.add(isotopeComponent!!) // element must be mode of isotopes
+                    }
+                } else {
+                    val isotopeComponent = this.materials.get<GdmlIsotope>(fraction.ref)
+                    if (isotopeComponent != null) {
+                        isotopesInUse.add(isotopeComponent)
                     }
                 }
             }
+        }
+    }
 
-            gdml.structure.content.forEach {
-                val volume = gdml.structure.get<GdmlVolume>(it.name)
-                if (volume != null) {
-                    val material = volume.materialref.resolve(gdml)!!
-                    materialsInUse.add(material)
-                    val composite = gdml.materials.get<GdmlComposite>(material.name)
-                    addMaterialsRecursively(composite)
-                }
-            }
+    this.structure.content.forEach {
+        val volume = this.structure.get<GdmlVolume>(it.name)
+        if (volume != null) {
+            val material = volume.materialref.resolve(this)!!
+            materialsInUse.add(material)
+            val composite = this.materials.get<GdmlComposite>(material.name)
+            addMaterialsRecursively(composite)
+        }
+    }
 
 
-            // materialsInUse now contains all materials in use
-            /*
-            println(" - ISOTOPES:")
-            isotopesInUse.forEach { println(it) }
-            println(" - ELEMENTS:")
-            elementsInUse.forEach { println(it) }
-            println(" - MATERIALS:")
-            materialsInUse.forEach { println(it) }
-            */
+    // materialsInUse now contains all materials in use
+    /*
+    println(" - ISOTOPES:")
+    isotopesInUse.forEach { println(it) }
+    println(" - ELEMENTS:")
+    elementsInUse.forEach { println(it) }
+    println(" - MATERIALS:")
+    materialsInUse.forEach { println(it) }
+    */
 
-            gdml.materials.content.clear()
-            // Refill materials only with used ones
-            isotopesInUse.forEach {
-                gdml.materials.content.add(it)
-            }
-            elementsInUse.forEach {
-                gdml.materials.content.add(it)
-            }
-            materialsInUse.forEach {
-                gdml.materials.content.add(it)
-            }
-
-            return gdml
-        }
+    this.materials.content.clear()
+    // Refill materials only with used ones
+    isotopesInUse.forEach {
+        this.materials.content.add(it)
     }
-}
\ No newline at end of file
+    elementsInUse.forEach {
+        this.materials.content.add(it)
+    }
+    materialsInUse.forEach {
+        this.materials.content.add(it)
+    }
+
+    return this
+}
diff --git a/src/jvmTest/kotlin/space/kscience/gdml/PredefinedGeometryTest.kt b/src/jvmTest/kotlin/space/kscience/gdml/PredefinedGeometryTest.kt
index d46b13f..e6de40b 100644
--- a/src/jvmTest/kotlin/space/kscience/gdml/PredefinedGeometryTest.kt
+++ b/src/jvmTest/kotlin/space/kscience/gdml/PredefinedGeometryTest.kt
@@ -3,9 +3,7 @@ package space.kscience.gdml
 import org.junit.jupiter.api.Test
 import java.io.File
 import kotlin.test.assertNotNull
-import kotlin.test.assertNull
 import kotlin.test.assertEquals
-import kotlin.test.assertNotEquals
 
 class PredefinedGeometryTest {
 
@@ -64,13 +62,14 @@ class PredefinedGeometryTest {
     @Test
     fun testRemoveUnusedMaterials() {
         val file = File("gdml-source/babyIAXO.gdml")
-        var gdml = Gdml.decodeFromStream(file.inputStream(), true)
-        assertNotNull(gdml.materials.get<GdmlMaterial>("G4_WATER"))
-        val gdmlAfter = GdmlMaterialPostProcessor.removeUnusedMaterials(gdml)
-        // assertNull(gdmlAfter.materials.get<GdmlMaterial>("G4_WATER")) TODO: shouldn't this be null? its not on the gdml
-        // assertNotEquals(gdmlAfter, gdml) TODO: not sure why this doesn't work either
-        val gdmlAfterAgain = GdmlMaterialPostProcessor.removeUnusedMaterials(gdmlAfter)
-        assertEquals(gdmlAfter, gdmlAfterAgain)
+        val gdml = Gdml.decodeFromStream(file.inputStream(), true)
+        val materialToRemove = gdml.materials.get<GdmlMaterial>("G4_WATER")
+        assertNotNull(materialToRemove)
+        assert(gdml.materials.content.contains(materialToRemove))
+        val gdmlMaterialsRemoved = gdml.removeUnusedMaterials()
+        assert(!gdmlMaterialsRemoved.materials.content.contains(materialToRemove))
+        val gdmlMaterialsRemovedAgain = gdmlMaterialsRemoved.removeUnusedMaterials()
+        assertEquals(gdmlMaterialsRemoved, gdmlMaterialsRemovedAgain)
 
         println(gdml)
     }