Minor improvements + potential bug fix.