Rewrite previous patch with cleaner code

Non functional change