The ChaCha20 algorithm is simple enough that if you're curious about how it works you could just write your own implementation (but don't use it in production!). RFC 7539RFC 7539 is a simple reference that also comes with examples so you can easily check for correctness as you go along. But the way it works is that you XOR the message with a keystream that is generated as a set of 64-byte blocks, where each block is "named" by three arguments:
- The key (secret);
- The nonce (not secret);
- The block counter (not secret).
The block is initialized with these values and some constants in the way that your code snippet shows. Then the block is scrambled with the doubleround operation ten times (ten double rounds = twenty rounds, that's why there's a "20" in the name). And now, critically, the initial state of the block is added into the result of that, and this step requires the key again—so if you erased it you wouldn't be able to carry out this final step.
After that final addition the key should not be easily recoverable from the resulting block (or otherwise the cipher is completely broken!). But nevertheless the block's content is a sensitive secret, because an attacker who sees it can decrypt the corresponding ciphertext block.
All you've got at this point, however, is one 64-byte block, which is good only for encrypting or decrypting one 64-byte block of the message. Unless 64 bytes is the largest message length you'll ever encrypt, you're going to need the key again and again to generate additional blocks. So you can only erase the key and blocks after you've encrypted or decrypted the whole message.
Since both the key and the keystream blocks are secret data, you'd want to use your operating system's facilities to prevent that memory from being swapped out to disk. And you'll want to safely zero out that memory after you're done with it.
Similar considerations apply to other ciphers, but the details are different. For example, many ciphers (but not ChaCha20) have a key schedule, i.e. a sequence of "subkeys" that are derived from the secret key. Often it is the case that after you've generated the key schedule you don't need the key any more, so you could zero out the key at that point, but the subkeys need to be protected, and you need to hang on to them because you reuse them for each block that you process.